app.py

from transformers import GPT2LMHeadModel, GPT2Tokenizer

def generate_text(prompt, model_name="gpt2-medium", max_length=100, num_return_sequences=1):
    # Charger le modèle et le tokenizer
    tokenizer = GPT2Tokenizer.from_pretrained('rombodawg/test_dataset_Codellama-3-8B')
    model = GPT2LMHeadModel.from_pretrained('rombodawg/test_dataset_Codellama-3-8B')

    # Encoder le prompt en tokens
    input_ids = tokenizer.encode(prompt, return_tensors="pt")

    # Générer du texte
    output = model.generate(
        input_ids,
        max_length=max_length,
        num_return_sequences=num_return_sequences,
        no_repeat_ngram_size=2,
        top_k=50,
        top_p=0.95,
        temperature=0.7,
        pad_token_id=tokenizer.eos_token_id
    )

    # Décoder les résultats en texte
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)

    return generated_text

# Exemple d'utilisation
prompt = "Une fois dans une forêt sombre, un aventurier solitaire"
generated_text = generate_text(prompt)
print(generated_text)