feat: add BART

app.py

@@ -16,23 +16,22 @@ logging.basicConfig(level=logging.INFO)
 DATA_DIR = "/data" if os.path.exists("/data") else "."
 DATASET_DIR = os.path.join(DATA_DIR, "rag_dataset")
 DATASET_PATH = os.path.join(DATASET_DIR, "dataset")
-MODEL_PATH = "facebook/bart-large-cnn"  # Changed to BART model which is better for summarization
+MODEL_PATH = "t5-small"  # Changed to T5-small for better CPU compatibility
 
 @st.cache_resource
 def load_local_model():
     """Load the local Hugging Face model"""
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
-    model = AutoModelForSeq2SeqLM.from_pretrained(
-        MODEL_PATH,
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=True,
-        device_map=None  # Let PyTorch handle device placement
-    )
-    
-    # Move model to CPU explicitly
-    model = model.cpu()
-    
-    return model, tokenizer
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+        model = AutoModelForSeq2SeqLM.from_pretrained(
+            MODEL_PATH,
+            device_map={"": "cpu"},  # Force CPU
+            torch_dtype=torch.float32
+        )
+        return model, tokenizer
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        return None, None
 
 def fetch_arxiv_papers(query, max_results=5):
     """Fetch papers from arXiv"""
@@ -144,6 +143,9 @@ def generate_answer(question, context, max_length=512):
     """Generate a comprehensive answer using the local model"""
     model, tokenizer = load_local_model()
     
+    if model is None or tokenizer is None:
+        return "Error: Could not load the model. Please try again later."
+    
     # Format the context as a structured query
     prompt = f"""Summarize the following research about autism and answer the question.
 
@@ -152,7 +154,7 @@ Research Context:
 
 Question: {question}
 
-Provide a detailed answer that includes:
+Instructions: Based on the research context above, provide a comprehensive answer that covers:
 1. Main findings from the research
 2. Research methods used
 3. Clinical implications
@@ -160,41 +162,42 @@ Provide a detailed answer that includes:
 
 If the research doesn't address the question directly, explain what information is missing."""
     
-    # Generate response
-    inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
-    
-    # Move inputs to the same device as model
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
-    
-    with torch.inference_mode():
-        outputs = model.generate(
-            **inputs,
-            max_length=max_length,
-            min_length=200,  # Ensure longer responses
-            num_beams=5,
-            length_penalty=2.0,  # Encourage even longer responses
-            temperature=0.7,
-            no_repeat_ngram_size=3,
-            repetition_penalty=1.3,
-            early_stopping=True
-        )
-    
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # If response is too short or empty, provide a fallback message
-    if len(response.strip()) < 100:
-        return """I apologize, but I couldn't generate a specific answer from the research papers provided. 
-        This might be because:
-        1. The research papers don't directly address your question
-        2. The context needs more specific information
-        3. The question might need to be more specific
+    try:
+        # Generate response
+        inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
         
-        Please try rephrasing your question or ask about a more specific aspect of autism."""
-    
-    # Format the response for better readability
-    formatted_response = response.replace(". ", ".\n").replace("• ", "\n• ")
+        with torch.inference_mode():
+            outputs = model.generate(
+                **inputs,
+                max_length=max_length,
+                min_length=100,
+                num_beams=4,
+                length_penalty=1.5,
+                temperature=0.7,
+                repetition_penalty=1.2,
+                early_stopping=True
+            )
+        
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # If response is too short or empty, provide a fallback message
+        if len(response.strip()) < 50:
+            return """I apologize, but I couldn't generate a specific answer from the research papers provided. 
+            This might be because:
+            1. The research papers don't directly address your question
+            2. The context needs more specific information
+            3. The question might need to be more specific
+            
+            Please try rephrasing your question or ask about a more specific aspect of autism."""
+        
+        # Format the response for better readability
+        formatted_response = response.replace(". ", ".\n").replace("• ", "\n• ")
+        
+        return formatted_response
     
-    return formatted_response
+    except Exception as e:
+        st.error(f"Error generating response: {str(e)}")
+        return "Error: Could not generate response. Please try again with a different question."
 
 # Streamlit App
 st.title("🧩 AMA Autism")

requirements.txt

@@ -4,10 +4,7 @@ datasets>=2.17.0
 --extra-index-url https://download.pytorch.org/whl/cpu
 torch>=2.2.0
 accelerate>=0.26.0
-safetensors>=0.4.1
 numpy>=1.24.0
 pandas>=2.2.0
 requests>=2.31.0
-arxiv>=2.1.0
-lancedb>=0.3.3
-tantivy>=0.19.2
+arxiv>=2.1.0