Update app.py

app.py

@@ -3,9 +3,7 @@ import gradio as gr
 import torch
 import pandas as pd
 from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
-from datasets import Dataset
 import yfinance as yf
-import numpy as np
 
 # Function to fetch and preprocess ICICI Bank data
 def fetch_and_preprocess_data():
@@ -26,51 +24,20 @@ def fetch_and_preprocess_data():
         print(f"Error fetching data: {e}")
         return pd.DataFrame()  # Return an empty DataFrame if fetching fails
 
-# Function to create and save a custom index for the retriever
-def create_custom_index():
-    # Fetch and preprocess data
-    data = fetch_and_preprocess_data()
-    
-    if data.empty:
-        raise ValueError("No data available to create the index.")
-    
-    # Create a dataset for the retriever
-    dataset = Dataset.from_dict({
-        "id": [str(i) for i in range(len(data))],
-        "text": data.apply(lambda row: f"Date: {row.name}, Close: {row['Close']:.2f}, MA_50: {row['MA_50']:.2f}, MA_200: {row['MA_200']:.2f}", axis=1).tolist(),
-        "title": [f"ICICI Bank Data {i}" for i in range(len(data))]
-    })
-    
-    # Save the dataset and index
-    dataset_path = "icici_bank_dataset"
-    index_path = "icici_bank_index"
-    dataset.save_to_disk(dataset_path)
-    print(f"Dataset saved to {dataset_path}")
-    
-    # Add FAISS index
-    dataset.add_faiss_index("text")
-    dataset.get_index("text").save(index_path)
-    print(f"FAISS index saved to {index_path}")
-    
-    return dataset_path, index_path
-
 # Load the fine-tuned RAG model and tokenizer
-tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-base")
-print("Tokenizer loaded successfully.")
-
 try:
-    # Create and save the custom index
-    dataset_path, index_path = create_custom_index()
-    
-    # Load the retriever with the custom index
+    tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-base")
+    print("Tokenizer loaded successfully.")
+
+    # Use a pre-built index (e.g., wiki_dpr) instead of creating a custom index
     retriever = RagRetriever.from_pretrained(
         "facebook/rag-sequence-base",
-        index_name="custom",
-        passages_path=dataset_path,
-        index_path=index_path
+        index_name="wiki_dpr",
+        passages_path=None,
+        index_path=None
     )
     print("Retriever loaded successfully.")
-    
+
     # Load the RAG model
     model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-base", retriever=retriever)
     print("Model loaded successfully.")