Data Extraction + Factuality Checks

1
import json
2
import itertools
3
4
import pandas as pd
5
from langchain import PromptTemplate
6
from predictionguard import PredictionGuard
7
8
# Set your Prediction Guard token as an environmental variable.
9
os.environ["PREDICTIONGUARD_API_KEY"] = "<api key>"
10
11
client = PredictionGuard()
12
13
# Load the JSON data into a dataframe
14
data = []
15
with open('transcripts.json') as f:
16
    for line in itertools.islice(f, 5):
17
        line = line.strip()
18
        if not line: continue
19
        data.append(json.loads(line))
20
df = pd.DataFrame(data)
21
22
# Transform rows to columns
23
df = df.transpose()
24
25
# Reset the index and assign an index column
26
df = df.reset_index()
27
df.columns = ['id', 'transcript']
28
29
# Reset the index (optional)
30
df.reset_index(drop=True, inplace=True)
31
32
# Start with 5 rows of the dataframe
33
df=df.head(5)

1
# Define the summarization prompt
2
summarize_template = """### Instruction:
3
Summarize the input transcript below.
4
5
### Input:
6
{transcript}
7
8
### Response:
9
"""
10
11
summary_prompt = PromptTemplate(template=summarize_template,
12
    input_variables=["context"],
13
)
14
15
# Loop over the rows summarizing the data
16
summaries = []
17
for i,row in df.iterrows():
18
    result=client.completions.create(
19
        model="Hermes-2-Pro-Llama-3-8B",
20
        prompt=summary_prompt.format(
21
            transcript=row['transcript']
22
        ),
23
        max_tokens=200,
24
        temperature=0.1
25
    )
26
    print(result['choices'][0]['text'])
27
    summaries.append(result['choices'][0]['text'])
28
29
# Print out some summaries to sanity check them
30
df['summary']=summaries
31
print(df.head(5))
32
33
# Save the resulting data
34
df.to_csv("summarized_transcripts.csv", index=False)

1
# Define the questions we want answered
2
questions=["What symptoms is the patient experiencing",
3
           "What is the Patient's name?",
4
           "When did the symptoms start?",
5
           "On a scale of 1 to 10, what level of pain is the patient experiencing?"]
6
7
# Define the question answering prompt
8
question_answer_template = """### Instruction:
9
Answer the following question {question} using the below doctor-patient transcript summary.
10
11
### Input:
12
{transcript_summary}
13
14
### Response:
15
"""
16
17
q_and_a_prompt = PromptTemplate(template=question_answer_template,
18
    input_variables=["question", "transcript_summary"],
19
)

1
answers = {q: [] for q in questions}
2
fact_scores = {q: [] for q in questions}
3
4
# Loop over the rows of the dataframe processing each transcript
5
for i, row in df.iterrows():
6
    for q in questions:
7
8
        # Extract the information
9
        result = client.completions.create(
10
            model="Hermes-2-Pro-Llama-3-8B",
11
            prompt=q_and_a_prompt.format(
12
                question=q, transcript_summary=row["summary"]
13
            ),
14
            max_tokens=200,
15
            temperature=0.1,
16
17
        )
18
19
        # Generate a factual consistency score
20
        fact_score =client.factuality.check(
21
            reference=row['summary'],
22
            text=result['choices'][0]['text']
23
        )
24
25
        # Append the results
26
        fact_scores[q].append(fact_score['checks'][0]['score'])
27
        answers[q].append(result["choices"][0]["text"])
28
29
# Add the answers and fact scores as new columns to the original DataFrame
30
for q in questions:
31
    df[f"{q}_answer"] = answers[q]
32
    df[f"{q}_fact_score"] = fact_scores[q]
33
34
# Show some results
35
print(df.head(2))
36
37
# Save the results
38
df.to_csv("answers_with_fact_scores.csv", index=False)