Update app.py

app.py

@@ -9,7 +9,7 @@ import markdown
 import cv2
 import numpy as np
 from PIL import Image
-from transformers import AutoProcessor, VisionEncoderDecoderModel
+from transformers import AutoProcessor, VisionEncoderDecoderModel, Gemma3nForConditionalGeneration, pipeline
 import torch
 import os
 import tempfile
@@ -312,10 +312,37 @@ except Exception as e:
     dolphin_model = None
     model_status = f"❌ Model failed to load: {str(e)}"
 
+# Initialize chatbot model
+try:
+    import os
+    # Get HuggingFace token from environment/secrets
+    hf_token = os.getenv('HF_TOKEN')
+    if hf_token:
+        os.environ['HF_TOKEN'] = hf_token
+    
+    chatbot_model = Gemma3nForConditionalGeneration.from_pretrained(
+        "google/gemma-3n-e4b-it",
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        use_auth_token=hf_token
+    ).eval()
+    
+    chatbot_processor = AutoProcessor.from_pretrained(
+        "google/gemma-3n-e4b-it",
+        use_auth_token=hf_token
+    )
+    
+    print("Chatbot model loaded successfully")
+except Exception as e:
+    print(f"Error loading chatbot model: {e}")
+    chatbot_model = None
+    chatbot_processor = None
+
 
 # Global state for managing tabs
 processed_markdown = ""
 show_results_tab = False
+chatbot_model = None
 
 
 def process_uploaded_pdf(pdf_file, progress=gr.Progress()):
@@ -484,7 +511,7 @@ with gr.Blocks(
                 send_btn = gr.Button("Send", variant="primary", scale=1)
             
             gr.Markdown(
-                "*Chat functionality will be implemented in the next version*",
+                "*Ask questions about your processed document. The AI will use the document content to provide accurate answers.*",
                 elem_id="chat-notice"
             )
     
@@ -513,12 +540,74 @@ with gr.Blocks(
         outputs=[chat_tab]
     )
     
-    # Placeholder chat functionality
-    def placeholder_chat(message, history):
-        return history + [["Coming soon: AI-powered document Q&A", "This feature will allow you to ask questions about your processed PDF document."]]
+    # Chatbot functionality
+    def chatbot_response(message, history):
+        if not message.strip():
+            return history
+        
+        if chatbot_model is None:
+            return history + [[message, "❌ Chatbot model not loaded. Please check your HuggingFace token."]]
+        
+        if not processed_markdown:
+            return history + [[message, "❌ Please process a PDF document first before asking questions."]]
+        
+        try:
+            # Create context with the processed document
+            context = f"Document content:\n{processed_markdown[:3000]}..." if len(processed_markdown) > 3000 else f"Document content:\n{processed_markdown}"
+            
+            # Create chat messages
+            messages = [
+                {
+                    "role": "system",
+                    "content": [{"type": "text", "text": "You are a helpful assistant that answers questions about documents. Use the provided document content to answer questions accurately."}]
+                },
+                {
+                    "role": "user",
+                    "content": [{"type": "text", "text": f"{context}\n\nQuestion: {message}"}]
+                }
+            ]
+            
+            # Process with the model
+            inputs = chatbot_processor.apply_chat_template(
+                messages,
+                add_generation_prompt=True,
+                tokenize=True,
+                return_dict=True,
+                return_tensors="pt",
+            ).to(chatbot_model.device)
+            
+            input_len = inputs["input_ids"].shape[-1]
+            
+            with torch.inference_mode():
+                generation = chatbot_model.generate(
+                    **inputs,
+                    max_new_tokens=300,
+                    do_sample=False,
+                    temperature=0.7,
+                    pad_token_id=chatbot_processor.tokenizer.pad_token_id
+                )
+                generation = generation[0][input_len:]
+            
+            response = chatbot_processor.decode(generation, skip_special_tokens=True)
+            
+            return history + [[message, response]]
+            
+        except Exception as e:
+            error_msg = f"❌ Error generating response: {str(e)}"
+            return history + [[message, error_msg]]
     
     send_btn.click(
-        fn=placeholder_chat,
+        fn=chatbot_response,
+        inputs=[msg_input, chatbot],
+        outputs=[chatbot]
+    ).then(
+        lambda: "",
+        outputs=[msg_input]
+    )
+    
+    # Also allow Enter key to send message
+    msg_input.submit(
+        fn=chatbot_response,
         inputs=[msg_input, chatbot],
         outputs=[chatbot]
     ).then(