kamyvision-api / scripts /train_fake_news.py
oyabun-dev's picture
deploy: 2026-04-02T00:05:48Z
55bcd2b
import os
import torch
import pandas as pd
from datasets import Dataset
from transformers import (
AutoTokenizer,
AutoModelForSequenceClassification,
TrainingArguments,
Trainer,
DataCollatorWithPadding
)
# 1. Configuration
# Pour le Français, on utilise un modèle multilingue (XLM-RoBERTa) ou dédié (CamemBERT)
MODEL_NAME = "xlm-roberta-base"
OUTPUT_DIR = "./results"
SAVE_DIR = "./fakenews_model_fr"
def train():
# 2. Chargement des données (Exemple avec un CSV fictif)
# Assurez-vous d'avoir un CSV avec au moins 'text' et 'label' (0=Real, 1=Fake)
if not os.path.exists("data.csv"):
print("Erreur: data.csv non trouvé. Création d'un exemple factice...")
data = {
"text": [
"Le président a inauguré une nouvelle école ce matin à Paris.",
"Des extraterrestres ont envahi la tour Eiffel hier soir !"
],
"label": [0, 1]
}
pd.DataFrame(data).to_csv("data.csv", index=False)
df = pd.read_csv("data.csv")
# Nettoyage des données : Supprimer les NaNs et s'assurer que 'text' est de type string
df = df.dropna(subset=['text', 'label'])
df['text'] = df['text'].astype(str)
df = df[df['text'].str.len() > 10]
# OPTIMISATION HACKATHON 🚀 : Échantillon aléatoire pour aller plus vite
MAX_SAMPLES = 2000
if len(df) > MAX_SAMPLES:
print(f"Extraction d'un échantillon aléatoire de {MAX_SAMPLES} sur {len(df)} lignes...")
df = df.sample(MAX_SAMPLES, random_state=42).reset_index(drop=True)
dataset = Dataset.from_pandas(df.reset_index(drop=True))
# 3. Tokenization
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
def preprocess_function(examples):
return tokenizer(examples["text"], truncation=True, padding=True)
tokenized_dataset = dataset.map(preprocess_function, batched=True)
# 4. Chargement du modèle
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)
# 5. Arguments d'entraînement
training_args = TrainingArguments(
output_dir=OUTPUT_DIR,
learning_rate=2e-5,
per_device_train_batch_size=8, # Réduit pour économiser la RAM
num_train_epochs=1, # 1 seule passe pour le hackathon
weight_decay=0.01,
save_strategy="epoch",
)
# 6. Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_dataset,
tokenizer=tokenizer,
data_collator=DataCollatorWithPadding(tokenizer=tokenizer),
)
# 7. Lancement de l'entraînement
print("Début de l'entraînement...")
trainer.train()
# 8. Sauvegarde
model.save_pretrained(SAVE_DIR)
tokenizer.save_pretrained(SAVE_DIR)
print(f"Modèle sauvegardé dans {SAVE_DIR}")
if __name__ == "__main__":
train()