Premier-gheya-innov

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import re
from typing import Optional

# --- Constantes (Doivent correspondre à celles utilisées dans l'entraînement) ---
# Tokens utilisés pour structurer l'entrée/sortie (à supprimer)
SEP_TOKEN = "<|sep|>"
PROMPT_STRUCTURE = "Question: "
RESPONSE_STRUCTURE = "Réponse:"

# Paramètres de génération
MAX_NEW_TOKENS = 300
TEMPERATURE = 0.7
MODEL_ID = "Clemylia/Premier-gheya-innov"

def generate_text_public(
    repo_id: str = MODEL_ID,
    prompt: str = "Test?",
    max_new_tokens: int = MAX_NEW_TOKENS,
    temperature: float = TEMPERATURE,
    device: Optional[str] = None
):
    """
    Charge le modèle Gheya-Innov et génère du texte avec un nettoyage strict de la sortie.
    """

    print(f"🔄 Chargement du Tokenizer et du Modèle depuis {repo_id}...")

    try:
        # Déterminer le dispositif (GPU si disponible, sinon CPU)
        if device is None:
            device = "cuda" if torch.cuda.is_available() else "cpu"
        device = torch.device(device)

        # Chargement du Tokenizer (standard)
        tokenizer = AutoTokenizer.from_pretrained(repo_id)

        # Chargement du Modèle (trust_remote_code=True est essentiel)
        model = AutoModelForCausalLM.from_pretrained(
            repo_id,
            trust_remote_code=True,
            torch_dtype=torch.float32, # Explicitly set dtype
            use_cache=False, # Disable cache to avoid TypeError with custom model
        ).to(device).eval()

    except Exception as e:
        print(f"❌ Erreur lors du chargement : {e}")
        print("\nATTENTION : Vérifiez que l'architecture et les poids sont correctement publiés sur le Hub.")
        return

    print(f"✅ Modèle chargé sur {device} et prêt pour la génération.")

    # 1. Préparation du Prompt
    # Format d'entrée : "Question: <question> <|sep|> Réponse:"
    formatted_prompt = f"{PROMPT_STRUCTURE}{prompt} {SEP_TOKEN} {RESPONSE_STRUCTURE}"
    print(f"\nPrompt formaté (pour le modèle): {formatted_prompt.replace(SEP_TOKEN, '[SEP]')}")

    # Encodage du prompt
    input_ids = tokenizer.encode(formatted_prompt, return_tensors='pt').to(device)

    # Crée l'attention_mask comme un LongTensor de 1 (pour contourner l'erreur de NoneType dans generate)
    attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=device)

    # 2. Génération
    print("⏳ Démarrage de la génération...")
    with torch.no_grad():
        output_ids = model.generate(
            input_ids,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            temperature=temperature,
            # Assure que la génération s'arrête si le modèle le demande
            pad_token_id=tokenizer.eos_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )

    # 3. Décodage
    # skip_special_tokens=False est nécessaire pour voir et nettoyer manuellement les balises
    generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=False)

    # 4. Nettoyage du Texte Généré (Rendement Propre)

    # Étape A: Isoler la réponse brute
    # On cherche le début de la réponse pour s'assurer de ne pas inclure le prompt
    try:
        start_index = generated_text.index(RESPONSE_STRUCTURE) + len(RESPONSE_STRUCTURE)
        response = generated_text[start_index:].strip()
    except ValueError:
        # Si la structure "Réponse:" n'est pas trouvée (très rare après entraînement), on prend tout
        response = generated_text

    # Étape B: Suppression des tokens spéciaux et des balises de formatage
    # 1. Suppression des tokens spéciaux (SEP_TOKEN, EOS, BOS, PAD)
    special_tokens = [SEP_TOKEN, tokenizer.eos_token, tokenizer.bos_token, tokenizer.pad_token]

    for token in special_tokens:
        if token is not None:
            # Utilisation de regex pour une suppression plus sûre et pour les cas où les tokens sont collés
            response = re.sub(re.escape(token), '', response, flags=re.IGNORECASE).strip()

    # 2. Suppression des balises de structure de l'input au cas où elles seraient répétées
    response = response.replace(PROMPT_STRUCTURE, "").strip()
    response = response.replace(RESPONSE_STRUCTURE, "").strip()
    
    # 3. Suppression spécifique du caractère 'Ġ' qui représente souvent un espace de préfixe
    response = response.replace('Ġ', '').strip()

    # Étape C: Suppression des répétitions, des sauts de ligne excessifs, et des espaces multiples
    response = re.sub(r'\s+', ' ', response).strip()

    # Si la réponse commence par le prompt d'origine, on le retire (garde en sécurité)
    if response.startswith(prompt):
        response = response[len(prompt):].strip()

    print("\n--- Résultat de Gheya-Innov ---")
    print(f"Question: {prompt}")
    print(f"Réponse: {response}")
    print("-------------------------------")


if __name__ == "__main__":
    # Test avec la question par défaut
    generate_text_public(prompt="Test?")