Update app.py

app.py

@@ -4,17 +4,18 @@ import fitz  # PyMuPDF
 import pytesseract
 from PIL import Image
 import gradio as gr
-from transformers import pipeline
+import os
 from extraction_service import ExtractionService
+from transformers import pipeline
 
 # Load field extraction config
 extractor = ExtractionService("fields_config.json")
 
-# Load the TinyLlama chat pipeline (text-generation with instruction-tuning)
+# Initialize Hugging Face text-generation pipeline with Dolly v2 3B
 chat_pipeline = pipeline(
     "text-generation",
-    model="TinyLlama/TinyLlama-1.1B-Chat",
-    device=-1,  # CPU; change to 0 if you have a GPU
+    model="databricks/dolly-v2-3b",
+    device=-1,  # CPU; set device=0 for GPU if available
     max_length=512,
     do_sample=True,
     temperature=0.7,
@@ -71,17 +72,21 @@ def ask_question(question):
         return "Upload and extract a PDF first."
 
     context = json.dumps(extracted_data_store["fields"], indent=2)
-    prompt = f"Context: {context}\n\nQuestion: {question}\nAnswer:"
-
+    prompt = f"You are an assistant for insurance data.\nExtracted data:\n{context}\n\nQuestion: {question}\nAnswer:"
+    
     try:
-        results = chat_pipeline(prompt)
-        answer = results[0]["generated_text"].split("Answer:")[-1].strip()
+        responses = chat_pipeline(prompt, max_length=512, do_sample=True, temperature=0.7, top_p=0.9)
+        answer = responses[0]["generated_text"]
+        
+        # Extract answer text after the question (remove prompt)
+        if prompt in answer:
+            answer = answer.split(prompt,1)[1].strip()
         return answer
     except Exception as e:
         return f"Model inference error: {str(e)}"
 
 with gr.Blocks() as demo:
-    gr.Markdown("## 🛡️ Insurance PDF Extractor & Q&A (TinyLlama 1.1B Chat)")
+    gr.Markdown("## 🛡️ Insurance PDF Extractor & Q&A using Dolly-v2-3B")
     with gr.Row():
         pdf_input = gr.File(label="Upload PDF", file_types=[".pdf"])
         extract_btn = gr.Button("Extract")