Update app.py

app.py

@@ -1,37 +1,44 @@
-from transformers import pipeline, AutoTokenizer, AutoModel
 import pandas as pd
-from sentence_transformers import SentenceTransformer, util
+from transformers import pipeline
+import gradio as gr
 
-# Load data
-data = pd.read_csv("Injury_History.csv")
+# Load the injury data
+injury_data = pd.read_csv("Injury_History.csv")
 
-# Load embeddings model for retrieval
-embedding_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
+# Initialize a text generation model (use 'gpt2' or similar for simplicity)
+generator = pipeline("text-generation", model="gpt2")
 
-# Load language model for generation
-generator = pipeline("text-generation", model="gpt2")  # Lightweight option
-
-
-# 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)
+# Define the RAG function for injury lookup and generation
+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)
     
-    # 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
+    # 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)]
     
-    # Step 4: Prepare context for generation
-    context = ". ".join(top_injuries['Notes'].values)
+    if filtered_data.empty:
+        return "No injury records found for this player in the specified timeframe."
     
-    # Step 5: Generate response
-    generated_response = generator(f"Based on the injury history: {context}", max_length=100)[0]["generated_text"]
+    # Concatenate injury records for context
+    injury_details = ". ".join(filtered_data['Notes'].tolist())
+    context = f"{player_name.capitalize()}'s injuries: {injury_details}"
     
-    return generated_response
+    # Generate a response
+    response = generator(f"Based on available data, here are the injuries: {context}", max_length=100)[0]['generated_text']
+    return response
+
+# Set up Gradio interface
+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?'"
+)
 
-# Example usage
-player_query = "Tell me about Jaylen Brown's injuries."
-print(RAG_injury_info(player_query))
+# Launch the app
+interface.launch()