|
|
|
|
|
from rag_pipeline import retrieve_context |
|
|
from transformers import pipeline |
|
|
|
|
|
|
|
|
|
|
|
summarizer = pipeline("summarization", model="facebook/bart-large-cnn") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def generate_summary(retrieved_text: str): |
|
|
|
|
|
result = summarizer(retrieved_text, max_length=250, min_length=80, do_sample=False) |
|
|
return result[0]['summary_text'] |
|
|
|
|
|
|
|
|
""" |
|
|
prompt = f''' |
|
|
You are a clinical summarization assistant. |
|
|
Use ONLY the provided context to create a structured summary. |
|
|
Do not invent information. |
|
|
|
|
|
Context: |
|
|
{retrieved_text} |
|
|
|
|
|
Write the output in this exact format: |
|
|
Chief Complaint: ... |
|
|
HPI: ... |
|
|
PMH: ... |
|
|
Medications: ... |
|
|
Allergies: ... |
|
|
Assessment: ... |
|
|
Plan: ... |
|
|
''' |
|
|
result = summarizer(prompt, max_new_tokens=300, do_sample=False) |
|
|
return result[0]['generated_text'] |
|
|
""" |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
query = "Summarize into HPI/Assessment/Plan" |
|
|
|
|
|
retrieved_text = retrieve_context(query, top_k=5) |
|
|
|
|
|
print("=== Retrieved Context ===") |
|
|
print(retrieved_text) |
|
|
print("\n=== Structured Clinical Summary ===") |
|
|
summary = generate_summary(retrieved_text) |
|
|
print(summary) |