adarshjha01's picture
Create app.py
c12a143 verified
import pandas as pd
from sklearn.linear_model import LinearRegression
import gradio as gr
# Load dataset
teams = pd.read_csv("teams.csv")
teams = teams[['team', 'country', 'year', 'athletes', 'age', 'prev_medals', 'medals']]
teams = teams.dropna()
# Split data into training and testing sets
train = teams[teams['year'] < 2012].copy()
test = teams[teams['year'] >= 2012].copy()
# Define predictors and target
predictors = ['athletes', 'prev_medals']
target = 'medals'
# Train the Linear Regression model
reg = LinearRegression()
reg.fit(train[predictors], train['medals'])
# Define the prediction function
def predict_medals(athletes: int, prev_medals: int):
input_data = pd.DataFrame({'athletes': [athletes], 'prev_medals': [prev_medals]})
prediction = reg.predict(input_data)[0]
return max(0, round(prediction))
# Create Gradio interface
interface = gr.Interface(
fn=predict_medals,
inputs=[
gr.Number(label="Number of Athletes"),
gr.Number(label="Previous Medals Won"),
],
outputs="number",
title="Olympics Medal Prediction",
description="Predict the number of medals a team might win based on athletes and previous medals."
)
# Launch the interface
if __name__ == "__main__":
interface.launch()