Nelya-v1

# ==============================================================================
# 1. INSTALLATION ET IMPORTS
# ==============================================================================
# Nous n'avons besoin que des bibliothèques de base pour l'inférence
!pip install -q transformers torch

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# ==============================================================================
# 2. DÉFINITION DES CONSTANTES ET CHARGEMENT
# ==============================================================================
MODEL_NAME = "Clemylia/Nelya-v1" # Votre modèle finetuné
# Utiliser "auto" pour charger le modèle sur le GPU si disponible, sinon sur CPU
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

print(f"Chargement du modèle {MODEL_NAME} sur {DEVICE}...")

# Chargement du tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Chargement du modèle. Utiliser torch_dtype=torch.bfloat16 si votre environnement
# (GPU T4) le supporte et que le modèle a été entraîné ou sauvegardé dans ce format.
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16 if DEVICE == "cuda" else None,
    device_map="auto"
)

# ==============================================================================
# 3. PRÉPARATION DE LA REQUÊTE ET GÉNÉRATION
# ==============================================================================

# La question simple génerique (nous utilisons le format Chat utilisé lors du fine-tuning)
QUESTION = "Ti eta ?"

# Formatage de la question dans le format de conversation attendu par le modèle
# <|user|>...</s><|assistant|>
prompt = f""### Instruction:{QUESTION}### Response:"

# Tokenization du prompt
inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)

# Génération de la réponse
print("\n--- Génération de la réponse ---")

# Paramètres de génération standard
output_tokens = model.generate(
    **inputs,
    max_new_tokens=150,      # Limite la longueur de la réponse
    do_sample=True,          # Utiliser l'échantillonnage pour une réponse plus créative
    temperature=0.7,         # Contrôle le caractère aléatoire
    top_k=50,
    eos_token_id=tokenizer.eos_token_id, # Arrêt à la fin de séquence
)

# ==============================================================================
# 4. DÉCODAGE ET AFFICHAGE DU RÉSULTAT
# ==============================================================================
# Décoder le résultat, en ignorant le prompt initial et le token de padding/spécial
generated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)

# Nettoyage pour n'afficher que la réponse de l'assistant
# On cherche le début de la réponse de l'assistant après le prompt formaté
assistant_prefix = "### Response:"
if assistant_prefix in generated_text:
    # Retirer le prompt initial et les balises de fin
    clean_response = generated_text.split(assistant_prefix)[-1].strip()
else:
    # Si le formatage n'a pas été parfait (ce qui arrive)
    clean_response = generated_text.replace(QUESTION, "").strip()


print(f"\n🧠 Réponse de {MODEL_NAME} :\n")
print(clean_response)
print("\n-------------------------------")

# ==============================================================================
# 5. UTILISATION OPTIONNELLE DE PIPELINE (plus simple)
# ==============================================================================
print("\nTest rapide avec le pipeline (méthode simplifiée) :")

generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer
    # device=0 if DEVICE == "cuda" else -1 # REMOVED: Conflicts with device_map="auto"
)

# Le pipeline gère souvent bien le formatage par défaut pour les modèles Chat/Instruction
# On utilise la même question, sans le formatage manuel
result = generator(
    QUESTION,
    max_new_tokens=150,
    do_sample=True,
    temperature=0.7
)
print(result[0]['generated_text'])