Update app.py

app.py

@@ -1,75 +1,97 @@
 import json
 import faiss
 from sentence_transformers import SentenceTransformer
-from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
-import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # Load the knowledge base
 with open("knowledge_base.json", "r") as file:
     kb = json.load(file)
 
+os.system("huggingface-cli login")
 # Initialize the embedding model
 embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
-
-# Generate embeddings for the knowledge base
 kb_texts = [f"{item['Component']} {item['Range']} {item['Advice']}" for item in kb]
 kb_embeddings = embedding_model.encode(kb_texts)
 
-# Create a FAISS index
+# Create FAISS index
 index = faiss.IndexFlatL2(kb_embeddings.shape[1])
 index.add(kb_embeddings)
 
-# Load Hugging Face LLM (flan-t5 model as an example)
-model_name = "google/flan-t5-large"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-llm = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-text_generator = pipeline("text2text-generation", model=llm, tokenizer=tokenizer)
+# Load the Hugging Face LLM
+llama_model_name = "meta-llama/Llama-3.2-3B-Instruct"
+API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
 
+tokenizer = AutoTokenizer.from_pretrained(llama_model_name, token=API_TOKEN)
+llm = AutoModelForCausalLM.from_pretrained(llama_model_name, token=API_TOKEN)
 
-# Function to generate advice using RAG
+# Generate advice using RAG
 def generate_advice(extracted_data):
     try:
         recommendations = []
         for item in extracted_data:
+            # Query FAISS for the closest match
             query = f"{item['Component']} {item['Status']}"
             query_embedding = embedding_model.encode([query])
-
-            # Retrieve nearest knowledge base entry
             _, idx = index.search(query_embedding, 1)
             best_match = kb[idx[0][0]]
 
-            # Use Hugging Face LLM to generate detailed advice
+            # Prepare the LLM prompt
+            role = "Medical expert providing advice based on lab results."
             prompt = f"""
             Lab Test: {item['Component']}
             Value: {item['Value']} {item['Units']}
             Status: {item['Status']}
             
-            Advice based on guidelines:
-            {best_match['Advice']}
+            Medical Guidelines: {best_match['Advice']}
+            
+            Provide additional insights or recommendations.
             """
-            response = text_generator(prompt, max_length=150, num_return_sequences=1)
+
+            message_yours = [
+                {
+                    "role": "system",
+                    "content": role,
+                },
+                {"role": "user", "content": prompt},
+            ]
+
+            input_text_with_your_role = tokenizer.apply_chat_template(
+                message_yours,
+                tokenize=False,
+                add_generation_prompt=True,
+                return_tensors="pt",
+            )
+
+            # Generate response
+            output = llm.generate(
+                **input_text_with_your_role,
+                max_length=150,
+                num_return_sequences=1
+            )
+            advice = tokenizer.decode(output[0], skip_special_tokens=True).strip()
+
             recommendations.append({
                 "Component": item["Component"],
-                "Advice": response[0]["generated_text"]
+                "Advice": advice
             })
         return recommendations
     except Exception as e:
         return f"Error: {e}"
 
+# Gradio app with LLM integration
+import gradio as gr
+from PyPDF2 import PdfReader
+import re
 
-# Function to extract structured data from the PDF text
+# Function to extract structured data from PDF
 def pdf_to_text(pdf_file):
     try:
-        from PyPDF2 import PdfReader
-
         reader = PdfReader(pdf_file.name)
         text = ""
         for page in reader.pages:
             text += page.extract_text()
 
         # Regex to match lab results (e.g., WBC 4.4 4.8 10.8 K/ul Low)
-        import re
-
         pattern = r"(\w+)\s+([\d.]+)\s+([\d.]+)\s+([\d.]+)\s+([\w/%]+)\s+(\w+)"
         matches = re.findall(pattern, text)
 
@@ -85,11 +107,10 @@ def pdf_to_text(pdf_file):
     except Exception as e:
         return f"Error: {e}"
 
-
-# Gradio Interface with Hugging Face LLM Integration
+# Gradio Interface
 def main():
     with gr.Blocks() as app:
-        gr.Markdown("## Medical Test Interpreter with RAG (Hugging Face)")
+        gr.Markdown("## Medical Test Interpreter with RAG and LLM")
         with gr.Row():
             pdf_input = gr.File(label="Upload PDF", type="filepath")
             structured_data = gr.JSON(label="Extracted Structured Data")
@@ -102,7 +123,6 @@ def main():
 
     app.launch()
 
-
 # Run the Gradio app
 if __name__ == "__main__":
     main()