src/train_baseline.py

import os
import pandas as pd
import joblib
from sklearn.model_selection import train_test_split
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, accuracy_score

def main():
    print("Chargement des données...")
    data_path = 'data/dataset_clean.csv'
    
    if not os.path.exists(data_path):
        print(f"Erreur : Le fichier {data_path} est introuvable. Veuillez d'abord exécuter data_preparation.py.")
        return

    df = pd.read_csv(data_path)
    
    # Vérification que les colonnes 'text' et 'label' existent
    if 'text' not in df.columns or 'label' not in df.columns:
        print("Erreur : Le dataset ne contient pas les colonnes 'text' et 'label' requises.")
        return

    # Gestion des valeurs nulles éventuelles
    df = df.dropna(subset=['text', 'label'])
    
    X = df['text']
    y = df['label']

    print("Séparation des données (80% Train, 20% Test)...")
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    print("Vectorisation du texte avec TfidfVectorizer (max_features=5000)...")
    vectorizer = TfidfVectorizer(max_features=5000)
    
    X_train_vec = vectorizer.fit_transform(X_train)
    X_test_vec = vectorizer.transform(X_test)

    print("Entraînement du modèle Logistic Regression de référence...")
    model = LogisticRegression(random_state=42, max_iter=1000)
    model.fit(X_train_vec, y_train)

    print("Évaluation du modèle sur le jeu de test...")
    y_pred = model.predict(X_test_vec)
    
    acc = accuracy_score(y_test, y_pred)
    report = classification_report(y_test, y_pred, target_names=["Humain (0)", "IA (1)"])

    print("\n" + "="*40)
    print("RÉSULTATS DE L'ÉVALUATION")
    print("="*40)
    print(f"Accuracy Score: {acc:.4f}\n")
    print("Classification Report:")
    print(report)
    print("="*40 + "\n")

    print("Sauvegarde du modèle et du vectoriseur...")
    os.makedirs('models', exist_ok=True)
    
    model_path = 'models/baseline_logreg_model.joblib'
    vectorizer_path = 'models/tfidf_vectorizer.joblib'
    
    joblib.dump(model, model_path)
    joblib.dump(vectorizer, vectorizer_path)
    
    print(f"Modèle sauvegardé dans : {model_path}")
    print(f"Vectoriseur sauvegardé dans : {vectorizer_path}")
    print("Terminé avec succès !")

if __name__ == "__main__":
    main()