Update app.py

app.py

@@ -1,27 +1,36 @@
+from transformers import pipeline, AutoTokenizer, AutoModel
 import pandas as pd
-import gradio as gr
+from sentence_transformers import SentenceTransformer, util
 
-# Load the injury data
+# Load data
 data = pd.read_csv("Injury_History.csv")
 
-# Define a function to look up injuries by player name
-def lookup_injuries(player_name):
-    # Filter injury data for the given player name
-    injuries = data[data['Name'].str.lower() == player_name.lower()]
-    if injuries.empty:
-        return f"No injury records found for {player_name}."
-    
-    # Format the output to show only Date and Notes columns
-    return injuries[['Date', 'Notes']].to_string(index=False)
+# Load embeddings model for retrieval
+embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+
+# Load language model for generation
+generator = pipeline("text-generation", model="gpt-3.5-turbo")  # Example model
 
-# Set up the Gradio interface
-interface = gr.Interface(
-    fn=lookup_injuries,
-    inputs="text",
-    outputs="text",
-    title="NBA Player Injury Lookup",
-    description="Enter a player's name to see their injury history."
-)
+# Step 1: Create embeddings of injury data
+data['embedding'] = data['Notes'].apply(lambda x: embedding_model.encode(x, convert_to_tensor=True))
+
+# Define RAG function
+def RAG_injury_info(player_query):
+    # Step 2: Create embedding for the user query
+    query_embedding = embedding_model.encode(player_query, convert_to_tensor=True)
+    
+    # Step 3: Compute cosine similarities
+    data['similarity'] = data['embedding'].apply(lambda x: util.cos_sim(query_embedding, x).item())
+    top_injuries = data.sort_values(by='similarity', ascending=False).head(3)  # Get top matches
+    
+    # Step 4: Prepare context for generation
+    context = ". ".join(top_injuries['Notes'].values)
+    
+    # Step 5: Generate response
+    generated_response = generator(f"Based on the injury history: {context}", max_length=100)[0]["generated_text"]
+    
+    return generated_response
 
-# Launch the app
-interface.launch()
+# Example usage
+player_query = "Tell me about Jaylen Brown's injuries."
+print(RAG_injury_info(player_query))