fix: dataframes

app.py

@@ -1,9 +1,10 @@
 import streamlit as st
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import os
-from datasets import load_from_disk
+from datasets import load_from_disk, Dataset
 import torch
 import logging
+import pandas as pd
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -29,7 +30,13 @@ def load_dataset():
             import faiss_index.index as idx
             papers = idx.fetch_arxiv_papers("autism research", max_results=100)
             idx.build_faiss_index(papers, dataset_dir=DATASET_DIR)
-    return load_from_disk(DATASET_PATH)
+    
+    # Load and convert to pandas for easier handling
+    dataset = load_from_disk(DATASET_PATH)
+    return pd.DataFrame({
+        'title': dataset['title'],
+        'text': dataset['text']
+    })
 
 def generate_answer(question, context, max_length=200):
     tokenizer, model = load_models()
@@ -46,12 +53,16 @@ def generate_answer(question, context, max_length=200):
     
     # Get model predictions
     with torch.no_grad():
-        outputs = model(**inputs)
-        answer_ids = torch.argmax(outputs.logits, dim=-1)
-        
-        # Convert token positions to text
-        answer = tokenizer.decode(answer_ids[0], skip_special_tokens=True)
-        
+        outputs = model.generate(
+            inputs["input_ids"],
+            max_length=max_length,
+            min_length=30,
+            num_beams=4,
+            length_penalty=2.0,
+            early_stopping=True
+        )
+        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
     return answer if answer and not answer.isspace() else "I cannot find a specific answer to this question in the provided context."
 
 # Streamlit App
@@ -61,12 +72,11 @@ query = st.text_input("Please ask me anything about autism ✨")
 if query:
     with st.status("Searching for answers..."):
         # Load dataset
-        dataset = load_dataset()
+        df = load_dataset()
         
         # Get relevant context
         context = "\n".join([
-            f"{paper['text'][:1000]}"  # Use more context for better answers
-            for paper in dataset[:3]
+            f"{text[:1000]}" for text in df['text'].head(3)
         ])
         
         # Generate answer
@@ -77,9 +87,9 @@ if query:
             st.write(answer)
             
             st.write("### Sources Used:")
-            for i in range(min(3, len(dataset))):
-                st.write(f"**Title:** {dataset[i]['title']}")
-                st.write(f"**Summary:** {dataset[i]['text'][:200]}...")
+            for _, row in df.head(3).iterrows():
+                st.write(f"**Title:** {row['title']}")
+                st.write(f"**Summary:** {row['text'][:200]}...")
                 st.write("---")
         else:
             st.warning("I couldn't find a specific answer in the research papers. Try rephrasing your question.")