Update app.py

app.py

@@ -2,13 +2,10 @@
 import gradio as gr
 import torch
 import pandas as pd
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration, RagConfig
+from datasets import Dataset
 import yfinance as yf
-
-# Load the fine-tuned RAG model and tokenizer
-tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-base")
-retriever = RagRetriever.from_pretrained("facebook/rag-sequence-base", index_name="custom")
-model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-base", retriever=retriever)
+import numpy as np
 
 # Function to fetch and preprocess ICICI Bank data
 def fetch_and_preprocess_data():
@@ -22,6 +19,44 @@ def fetch_and_preprocess_data():
     
     return data
 
+# Function to create and save a custom index for the retriever
+def create_custom_index():
+    # Fetch and preprocess data
+    data = fetch_and_preprocess_data()
+    
+    # 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)
+    dataset.add_faiss_index("text")
+    dataset.get_index("text").save(index_path)
+    
+    return dataset_path, index_path
+
+# Load the fine-tuned RAG model and tokenizer
+tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-base")
+
+# Create and save the custom index
+dataset_path, index_path = create_custom_index()
+
+# Load the retriever with the custom index
+retriever = RagRetriever.from_pretrained(
+    "facebook/rag-sequence-base",
+    index_name="custom",
+    passages_path=dataset_path,
+    index_path=index_path
+)
+
+# Load the RAG model
+model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-base", retriever=retriever)
+
 # Function to analyze trading data using the RAG model
 def analyze_trading_data(question):
     # Fetch and preprocess data