Update app.py

app.py

@@ -1,35 +1,36 @@
-# app.py
 import gradio as gr
 import torch
 import pandas as pd
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
 import yfinance as yf
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+
+# Check if GPU is available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
 
 # Function to fetch and preprocess ICICI Bank data
 def fetch_and_preprocess_data():
     try:
-        # Fetch ICICI Bank data using yfinance
-        ticker = "ICICIBANK.BO"  # Use BSE symbol
+        ticker = "ICICIBANK.BO"  # ICICI Bank BSE Symbol
         data = yf.download(ticker, start="2020-01-01", end="2023-01-01")
         
         if data.empty:
             raise ValueError("No data found for the given symbol.")
         
-        # Calculate technical indicators
+        # Calculate Moving Averages
         data['MA_50'] = data['Close'].rolling(window=50).mean()
         data['MA_200'] = data['Close'].rolling(window=200).mean()
         
         return data
     except Exception as e:
         print(f"Error fetching data: {e}")
-        return pd.DataFrame()  # Return an empty DataFrame if fetching fails
+        return pd.DataFrame()  # Return empty DataFrame if fetching fails
 
-# Load the fine-tuned RAG model and tokenizer
+# Load the RAG model and tokenizer with error handling
 try:
     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="wiki_dpr",
@@ -38,45 +39,44 @@ try:
     )
     print("Retriever loaded successfully.")
 
-    # Load the RAG model
-    model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-base", retriever=retriever)
+    model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-base", retriever=retriever).to(device)
     print("Model loaded successfully.")
 except Exception as e:
-    print(f"Error initializing model or retriever: {e}")
+    print(f"Error loading model or retriever: {e}")
     retriever = None
     model = None
 
-# Function to analyze trading data using the RAG model
+# Function to analyze trading data
 def analyze_trading_data(question):
     if model is None or retriever is None:
         return "Error: Model or retriever is not initialized. Please check the logs."
-    
+
     # Fetch and preprocess data
     data = fetch_and_preprocess_data()
-    
+
     if data.empty:
         return "Error: No data available for analysis."
-    
-    # Prepare context for the RAG model
+
+    # Prepare context for RAG model
     context = (
         f"ICICI Bank stock data:\n"
         f"Latest Close Price: {data['Close'].iloc[-1]:.2f}\n"
         f"50-Day Moving Average: {data['MA_50'].iloc[-1]:.2f}\n"
         f"200-Day Moving Average: {data['MA_200'].iloc[-1]:.2f}\n"
     )
-    
+
     # Combine question and context
     input_text = f"Question: {question}\nContext: {context}"
-    
-    # Tokenize the input
-    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
-    
-    # Generate the answer using the RAG model
+
+    # Tokenize input
+    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True).to(device)
+
+    # Generate answer using the model
     outputs = model.generate(inputs['input_ids'])
     
-    # Decode the output to get the answer
+    # Decode output
     answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
+
     return answer
 
 # Gradio interface
@@ -94,4 +94,5 @@ iface = gr.Interface(
 )
 
 # Launch the app
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()