Spaces:
Sleeping
Sleeping
File size: 3,948 Bytes
abdf1bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
# model/train_model.py
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.preprocessing import LabelEncoder
import joblib
import os
import logging
# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Load datasets
try:
matches_df = pd.read_csv('data/results.csv')
except FileNotFoundError as e:
logger.error(f"Dataset not found: {e}")
raise
matches_df['home_score'] = pd.to_numeric(matches_df['home_score'], errors='coerce').fillna(0)
matches_df['away_score'] = pd.to_numeric(matches_df['away_score'], errors='coerce').fillna(0)
# Define the training function
def train_and_save_models():
# --- Prepare Data ---
# Create a symmetric outcome: 0 for team1 win, 1 for draw, 2 for team2 win
def get_match_outcome(row):
if row['home_score'] > row['away_score']:
return 0 if row['home_team'] < row['away_team'] else 2
elif row['home_score'] < row['away_score']:
return 2 if row['home_team'] < row['away_team'] else 0
else:
return 1
matches_df['outcome'] = matches_df.apply(get_match_outcome, axis=1)
# Sort teams alphabetically to ensure consistency
matches_df['team1'] = matches_df.apply(lambda x: min(x['home_team'], x['away_team']), axis=1)
matches_df['team2'] = matches_df.apply(lambda x: max(x['home_team'], x['away_team']), axis=1)
# Get all unique team names from both home_team and away_team
all_teams = pd.concat([matches_df['home_team'], matches_df['away_team']]).unique()
# Encode team names with a single LabelEncoder fitted on all teams
le_outcome = LabelEncoder()
le_outcome.fit(all_teams) # Fit on all unique teams
# --- Logistic Regression for Match Outcome ---
X_outcome = pd.DataFrame({
'team1': le_outcome.transform(matches_df['team1']),
'team2': le_outcome.transform(matches_df['team2'])
})
y_outcome = matches_df['outcome']
# Split data and train Logistic Regression model
X_train_outcome, _, y_train_outcome, _ = train_test_split(X_outcome, y_outcome, test_size=0.2, random_state=42)
logistic_model = LogisticRegression(multi_class='multinomial', max_iter=1000)
logistic_model.fit(X_train_outcome, y_train_outcome)
# --- Linear Regression for Goal Prediction ---
X_goals = pd.DataFrame({
'team1': le_outcome.transform(matches_df['team1']),
'team2': le_outcome.transform(matches_df['team2'])
})
# Targets: home_score and away_score as separate predictions
y_team1_goals = matches_df['home_score'] # Goals scored by team1 (home team in original data)
y_team2_goals = matches_df['away_score'] # Goals scored by team2 (away team in original data)
# Split data for goal prediction
X_train_goals, _, y_train_team1_goals, _, y_train_team2_goals, _ = train_test_split(
X_goals, y_team1_goals, y_team2_goals, test_size=0.2, random_state=42
)
# Train Linear Regression models for each team's goals
linear_model_team1 = LinearRegression()
linear_model_team2 = LinearRegression()
linear_model_team1.fit(X_train_goals, y_train_team1_goals)
linear_model_team2.fit(X_train_goals, y_train_team2_goals)
# Ensure the model directory exists
os.makedirs('model', exist_ok=True)
# Save all models and the label encoder
joblib.dump(logistic_model, 'model/logistic_regression_model.pkl')
joblib.dump(linear_model_team1, 'model/linear_regression_team1_goals.pkl')
joblib.dump(linear_model_team2, 'model/linear_regression_team2_goals.pkl')
joblib.dump(le_outcome, 'model/label_encoder.pkl')
logger.info("Logistic Regression and Linear Regression models, along with LabelEncoder, saved successfully.")
if __name__ == "__main__":
train_and_save_models() |