Update app.py

app.py

@@ -1,16 +1,20 @@
 from datasets import load_dataset
 import gradio as gr
 from transformers import RagTokenizer, RagRetriever, RagTokenForGeneration
+import torch
 
 # Load the horoscope-chat dataset
 dataset = load_dataset("chloeliu/horoscope-chat", split="train")
 
-# Convert the dataset to a format that can be used by the retriever
+# Check dataset structure
+print(dataset.column_names)
+
+# Assuming 'input' and 'response' are the correct keys:
 def prepare_docs(dataset):
     docs = []
     for data in dataset:
-        question = data['question']
-        answer = data['response']
+        question = data.get('input', '')  # Safely access the 'input' field
+        answer = data.get('response', '')  # Safely access the 'response' field
         docs.append({
             "question": question,
             "answer": answer