Update app.py

app.py

@@ -12,6 +12,7 @@ import faiss
 import numpy as np
 import tempfile
 from PIL import Image
+from transformers import BitsAndBytesConfig
 import logging
 
 # Set up logging
@@ -50,15 +51,24 @@ def initialize_models():
             "question-answering",
             model="distilbert-base-cased-distilled-squad",
             tokenizer="distilbert-base-cased",
-            device=0 if device == "cuda" else -1
+            device=-1  # Force CPU for free tier
         )
         
         logger.info("Loading language model...")
-        tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+        model_name = "Qwen/Qwen2.5-1.5B-Instruct"  # Replace distilgpt2
+        # Configure 4-bit quantization
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True
+        )
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForCausalLM.from_pretrained(
-            "distilgpt2",
+            model_name,
+            quantization_config=quantization_config,  # Use 4-bit
             device_map="auto",
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            torch_dtype=torch.float16  # Optimize for CPU fallback
         )
         
         if tokenizer.pad_token is None:
@@ -70,6 +80,87 @@ def initialize_models():
         logger.error(f"Error initializing models: {str(e)}")
         raise
 
+# Generation-based answering
+def answer_with_generation(index, embeddings, chunks, question):
+    try:
+        logger.info(f"Answering with generation model: '{question}'")
+        global tokenizer, model
+        
+        if tokenizer is None or model is None:
+            logger.info("Generation models not initialized, creating now...")
+            model_name = "Qwen/Qwen2.5-1.5B-Instruct"
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True
+            )
+            tokenizer = AutoTokenizer.from_pretrained(model_name)
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                quantization_config=quantization_config,
+                device_map="auto",
+                torch_dtype=torch.float16
+            )
+            
+            if tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
+                model.config.pad_token_id = model.config.eos_token_id
+        
+        # Get embeddings for question
+        q_embedding = embedder.encode([question])
+        
+        # Find relevant chunks
+        _, top_k_indices = index.search(q_embedding, k=3)
+        relevant_chunks = [chunks[i] for i in top_k_indices[0]]
+        context = " ".join(relevant_chunks)
+        
+        # Limit context size for efficiency
+        if len(context) > 2000:  # Reduced for Qwen's efficiency
+            context = context[:2000]
+
+        # Create prompt (optimized for Qwen's instruction format)
+        prompt = f"""<|im_start|>system
+You are a helpful assistant answering questions based on provided PDF content. Use the information below to give a clear, concise, and accurate answer. Avoid speculation and focus on the context.
+<|im_end|>
+<|im_start|>user
+**Context**: {context}
+
+**Question**: {question}
+
+**Instruction**: Provide a detailed and accurate answer based on the context. If the context doesn't contain enough information, say so clearly. <|im_end|>"""
+        
+        # Handle inputs
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)  # Increased for Qwen
+        
+        # Move inputs to CPU (free tier)
+        inputs = {k: v.to('cpu') for k, v in inputs.items()}
+        
+        # Generate answer
+        output = model.generate(
+            **inputs,
+            max_new_tokens=300,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            num_beams=2,  # Reduced for speed
+            no_repeat_ngram_size=2
+        )
+        
+        # Decode and format answer
+        answer = tokenizer.decode(output[0], skip_special_tokens=True)
+        # Extract the answer after the instruction
+        if "<|im_end|>" in answer:
+            answer = answer.split("<|im_end|>")[1].strip()
+        elif "Instruction" in answer:
+            answer = answer.split("Instruction")[1].strip()
+        
+        logger.info(f"Generation answer: '{answer[:50]}...' (length: {len(answer)})")
+        return answer.strip()
+    except Exception as e:
+        logger.error(f"Generation error: {str(e)}")
+        return "I couldn't generate a good answer based on the PDF content."
+
 # Cleanup function for temporary files
 def cleanup_temp_files(filepath):
     try: