import gradio as gr
from transformers import GPTNeoForCausalLM, GPT2Tokenizer, Trainer, TrainingArguments
from datasets import load_dataset

# Charger le modèle GPT-Neo et le tokenizer
model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")

# Ajouter un token de padding
tokenizer.pad_token = tokenizer.eos_token  # Utiliser le token de fin de séquence comme token de padding

# Charger les données (remplacer par le chemin vers vos propres données)
dataset = load_dataset("json", data_files={"train": "Databep.jsonl"})

# Tokeniser les données
def tokenize_function(examples):
    return tokenizer(examples["prompt"], padding="max_length", truncation=True)

# Tokenisation des données
tokenized_datasets = dataset.map(tokenize_function, batched=True)

# Arguments d'entraînement
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=3,               # Nombre d'époques
    per_device_train_batch_size=4,    # Taille du batch, ajustez selon vos ressources
    save_steps=10_000,                # Sauvegarder tous les 10 000 steps
    save_total_limit=2,               # Conserver seulement 2 checkpoints
)

# Initialiser le Trainer pour fine-tuner GPT-Neo
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
)

# Lancer le fine-tuning
trainer.train()

# Sauvegarder le modèle fine-tuné
model.save_pretrained("./fine_tuned_gpt_neo")
tokenizer.save_pretrained("./fine_tuned_gpt_neo")

# Interface Gradio pour tester le modèle fine-tuné
def generate_text(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(inputs["input_ids"], max_length=150)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Créer une interface avec Gradio pour interagir avec le modèle
interface = gr.Interface(fn=generate_text, inputs="text", outputs="text")
interface.launch()