Update app.py

app.py

@@ -1,20 +1,25 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoModelForSequenceClassification
 from PIL import Image
 import pytesseract
-import PyPDF2
 import pdfplumber
 import torch
 
-# Load ClinicalBERT model for medical text analysis
+# Load BART for zero-shot classification and Bio_ClinicalBERT for text summarization
 @st.cache_resource
-def load_clinicalbert():
-    tokenizer = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
-    model = AutoModelForSequenceClassification.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
-    return pipeline("zero-shot-classification", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
+def load_models():
+    # Bio_ClinicalBERT for text summarization
+    tokenizer_bert = AutoTokenizer.from_pretrained("emilyalsentzer/Bio_ClinicalBERT")
+    model_bert = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
+    summarizer = pipeline("summarization", model=model_bert, tokenizer=tokenizer_bert, device=0 if torch.cuda.is_available() else -1)
 
-# Initialize ClinicalBERT classifier
-clinical_bert = load_clinicalbert()
+    # BART model for zero-shot classification
+    classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli", device=0 if torch.cuda.is_available() else -1)
+
+    return summarizer, classifier
+
+summarizer, classifier = load_models()
 
 # OCR for Image using Tesseract
 def extract_text_from_image(image):
@@ -28,34 +33,33 @@ def extract_text_from_pdf(pdf_file):
             text += page.extract_text() or ""
     return text
 
-# Analyze and interpret the medical report using ClinicalBERT
+# Analyze and interpret the medical report
 def analyze_medical_text(text):
-    # Summarize the extracted text
-    summary = clinical_bert(text, candidate_labels=["summary", "overview", "findings"], multi_label=False)['labels'][0]
+    # Summarize the extracted text using ClinicalBERT
+    summarized_text = summarizer(text, max_length=100, min_length=30, do_sample=False)[0]['summary_text']
 
-    # Provide detailed interpretation
-    interpretation = clinical_bert(
-        text,
+    # Use BART for classification insights
+    interpretation = classifier(
+        summarized_text,
         candidate_labels=["normal", "abnormal", "urgent", "needs follow-up", "critical condition"],
         multi_label=True
     )
 
-    # Provide actionable recommendations
-    recommendations = clinical_bert(
-        text,
+    recommendations = classifier(
+        summarized_text,
         candidate_labels=["medication", "dietary change", "exercise", "follow-up with a doctor", "lifestyle change"],
         multi_label=True
     )
 
     return {
-        "summary": summary,
+        "summary": summarized_text,
         "interpretation": interpretation['labels'],
         "recommendations": recommendations['labels']
     }
 
 # Streamlit UI
-st.title("Medical Lab Report Analyzer with ClinicalBERT")
-st.write("Upload your medical lab report (PDF/Image) to get a summary and actionable insights using ClinicalBERT.")
+st.title("Medical Lab Report Analyzer with ClinicalBERT and BART")
+st.write("Upload your medical lab report (PDF/Image) to get a summary and actionable insights.")
 
 uploaded_file = st.file_uploader("Choose a PDF/Image file", type=["pdf", "png", "jpg", "jpeg"])