Update app.py

app.py

@@ -1,71 +1,28 @@
-import numpy as np
-from datasets import load_dataset
 import gradio as gr
 from transformers import RagTokenizer, RagRetriever, RagTokenForGeneration
-import torch
+import json
 
-# Load the horoscope-chat dataset
-dataset = load_dataset("chloeliu/horoscope-chat", split="train")
+# Load horoscope data
+with open("horoscope_data.json", "r") as file:
+    horoscope_data = json.load(file)
 
-# Assuming 'input' and 'response' are the correct keys:
-def prepare_docs(dataset):
-    docs = []
-    for data in dataset:
-        question = data.get('input', '')  # Safely access the 'input' field
-        answer = data.get('response', '')  # Safely access the 'response' field
-        docs.append({
-            "question": question,
-            "answer": answer
-        })
-    return docs
+# Custom Retriever that looks up horoscopes
+class CustomHoroscopeRetriever(RagRetriever):
+    def __init__(self, horoscope_data):
+        self.horoscope_data = horoscope_data
 
-# Prepare the documents
-docs = prepare_docs(dataset)
-
-# Custom Retriever that searches in the dataset
-class HoroscopeRetriever(RagRetriever):
-    def __init__(self, docs, tokenizer):
-        self.docs = docs
-        self.tokenizer = tokenizer
-
-    def retrieve(self, question_hidden_states, n_docs=1):
-        # Convert the question_hidden_states to a text string
-        question = question_hidden_states[0]
-        
-        if isinstance(question, np.ndarray):
-            if question.size == 1:
-                question = question.item()  # Convert single-element array to scalar
-            else:
-                question = str(question[0])  # Take the first element of the array
-        else:
-            question = str(question)
-        
-        question = question.lower()
-        
-        # Simple retrieval logic: find the most relevant document based on the question
-        best_match = None
-        for doc in self.docs:
-            if question in doc["question"].lower():
-                best_match = doc
-                break
-        
-        if best_match:
-            # Fake embedding as RAG expects this (In a real case, compute embeddings)
-            retrieved_doc_embeds = torch.zeros((1, 1, 768))  # Example tensor
-            doc_ids = ["0"]  # Example document ID
-            docs = [best_match["answer"]]
+    def retrieve(self, question_texts, n_docs=1):
+        zodiac_sign = question_texts[0].capitalize()
+        if zodiac_sign in self.horoscope_data:
+            return [self.horoscope_data[zodiac_sign]]
         else:
-            retrieved_doc_embeds = torch.zeros((1, 1, 768))  # Example tensor
-            doc_ids = ["0"]  # Example document ID
-            docs = ["Sorry, I couldn't find a relevant horoscope."]
-        
-        return retrieved_doc_embeds, doc_ids, docs
+            return ["I couldn't find your zodiac sign. Please try again with a valid one."]
 
-# Initialize the custom retriever with the dataset
-tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-base")
-retriever = HoroscopeRetriever(docs, tokenizer)
+# Initialize the custom retriever with the horoscope data
+retriever = CustomHoroscopeRetriever(horoscope_data)
 
 # Initialize RAG components
+tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-base")
 model = RagTokenForGeneration.from_pretrained("facebook/rag-token-base", retriever=retriever)
 
 # Define the chatbot function
@@ -78,5 +35,5 @@ def horoscope_chatbot(input_text):
 # Set up Gradio interface
 iface = gr.Interface(fn=horoscope_chatbot, inputs="text", outputs="text", title="Horoscope RAG Chatbot")
 
-# Launch the interface with sharing enabled
-iface.launch(share=True)
+# Launch the interface
+iface.launch()