Spaces:
Sleeping
Sleeping
| import pandas as pd | |
| import sqlite3 | |
| from sklearn.model_selection import train_test_split | |
| from sklearn.ensemble import RandomForestClassifier | |
| from sklearn.metrics import classification_report, accuracy_score | |
| from sklearn.preprocessing import LabelEncoder | |
| import joblib | |
| def train_model(db_name): | |
| conn = sqlite3.connect(db_name) | |
| df = pd.read_sql('SELECT * FROM sharks', conn) | |
| conn.close() | |
| # Select features: Activity and Month | |
| # We need to filter out 'Unknown' or 0 months if they are too many, | |
| # but for precision let's use what we have. | |
| data = df[['Activity', 'Month', 'is_fatal']].dropna() | |
| # Encode categorical Activity | |
| le = LabelEncoder() | |
| data['Activity_Encoded'] = le.fit_transform(data['Activity'].astype(str)) | |
| X = data[['Activity_Encoded', 'Month']] | |
| y = data['is_fatal'] | |
| # Split with stratification | |
| X_train, X_test, y_train, y_test = train_test_split( | |
| X, y, test_size=0.2, random_state=42, stratify=y | |
| ) | |
| print("Training Random Forest Classifier...") | |
| rf = RandomForestClassifier(n_estimators=100, random_state=42) | |
| rf.fit(X_train, y_train) | |
| # Predict | |
| y_pred = rf.predict(X_test) | |
| report = classification_report(y_test, y_pred) | |
| acc = accuracy_score(y_test, y_pred) | |
| print("\nModel Evaluation:") | |
| print(f"Accuracy: {acc:.4f}") | |
| print("\nClassification Report:") | |
| print(report) | |
| # Save results to file | |
| with open('model_results.txt', 'w') as f: | |
| f.write(f"Accuracy: {acc:.4f}\n") | |
| f.write("\nClassification Report:\n") | |
| f.write(report) | |
| # Save model | |
| joblib.dump(rf, 'fatality_predictor.pkl') | |
| joblib.dump(le, 'activity_encoder.pkl') | |
| print("Model and encoder saved and results written to model_results.txt.") | |
| if __name__ == "__main__": | |
| try: | |
| train_model('master_sharks.db') | |
| except Exception as e: | |
| print(f"Modeling failed: {e}") | |