Update app.py

app.py

@@ -1,44 +1,63 @@
 import pandas as pd
+import torch
+from sentence_transformers import SentenceTransformer
+import faiss
 from transformers import pipeline
 import gradio as gr
 
 # Load the injury data
 injury_data = pd.read_csv("Injury_History.csv")
 
-# Initialize a text generation model (use 'gpt2' or similar for simplicity)
+# Initialize an embedding model for creating embeddings of the injury descriptions
+embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
+
+# Generate embeddings for each injury record in the dataset
+injury_data['embedding'] = injury_data['Notes'].apply(lambda x: embedding_model.encode(x, convert_to_tensor=True))
+
+# Convert embeddings to numpy arrays for FAISS
+embeddings = torch.stack(injury_data['embedding'].to_list()).cpu().numpy()
+
+# Set up a FAISS index for efficient similarity search
+index = faiss.IndexFlatL2(embeddings.shape[1])
+index.add(embeddings)
+
+# Define a function to retrieve injuries based on similarity to the query
+def retrieve_injuries(query):
+    # Generate an embedding for the user query
+    query_embedding = embedding_model.encode(query, convert_to_tensor=True).cpu().numpy()
+
+    # Search the FAISS index for the top 3 similar injuries
+    k = 3  # number of results to retrieve
+    distances, indices = index.search(query_embedding.reshape(1, -1), k)
+    
+    # Retrieve the most relevant injury records
+    results = injury_data.iloc[indices[0]]
+    return results
+
+# Initialize a text generation model for generating responses
 generator = pipeline("text-generation", model="gpt2")
 
-# Define the RAG function for injury lookup and generation
+# Define the main function to handle the user query, retrieve relevant injuries, and generate a response
 def injury_query(player_query):
-    # Extract player name and optional year from the query
-    words = player_query.lower().split()
-    player_name = " ".join(words[:2])  # assuming the first two words are the player's name
-    year = next((word for word in words if word.isdigit()), None)
-    
-    # Filter the data by player name and year if provided
-    filtered_data = injury_data[injury_data['Name'].str.lower() == player_name]
-    if year:
-        filtered_data = filtered_data[filtered_data['Date'].str.contains(year)]
-    
-    if filtered_data.empty:
-        return "No injury records found for this player in the specified timeframe."
+    # Retrieve relevant injury data
+    retrieved_injuries = retrieve_injuries(player_query)
     
-    # Concatenate injury records for context
-    injury_details = ". ".join(filtered_data['Notes'].tolist())
-    context = f"{player_name.capitalize()}'s injuries: {injury_details}"
+    # Combine injury details into a context string for generation
+    injury_details = ". ".join(retrieved_injuries['Notes'].tolist())
+    context = f"Injury history: {injury_details}"
     
-    # Generate a response
-    response = generator(f"Based on available data, here are the injuries: {context}", max_length=100)[0]['generated_text']
+    # Generate a response based on the retrieved data
+    response = generator(f"Answer based on data: {context}", max_length=100)[0]['generated_text']
     return response
 
-# Set up Gradio interface
+# Set up the Gradio interface for the app
 interface = gr.Interface(
     fn=injury_query,
     inputs="text",
     outputs="text",
     title="NBA Player Injury Q&A",
-    description="Ask about a player's injury history, e.g., 'What injuries did Jaylen Brown have in 2017?'"
+    description="Ask about a player's injury history, or inquire about common injuries."
 )
 
-# Launch the app
+# Launch the Gradio app
 interface.launch()