Tiny-lamina

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# --- CONFIGURATION ---
# Utilise le nom de ton nouveau modèle publié
MODEL_REPO = "Clemylia/Tiny-lamina" 
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

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

# 1. Chargement du Tokenizer et du Modèle
# AutoModelForCausalLM est la classe standard pour les modèles de génération (comme Lamina)
tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
model = AutoModelForCausalLM.from_pretrained(MODEL_REPO).to(DEVICE)

print("Modèle chargé avec succès.")

# --- 2. FONCTION DE GÉNÉRATION ---
def generate_lamina_text(prompt, max_new_tokens=50, num_return_sequences=1):
    """
    Génère du texte à partir d'un prompt donné.

    Args:
        prompt (str): La phrase ou la question de départ.
        max_new_tokens (int): Le nombre maximum de tokens à générer.
    
    Returns:
        str: Le texte généré par le modèle.
    """
    # 2.1 Tokenisation du prompt
    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
    
    # 2.2 Génération des tokens
    # Utilisation de 'do_sample=True' pour une génération plus créative (comme le concept de Lamination)
    # Utilisation de 'top_k' ou 'top_p' pour contrôler la créativité
    generated_ids = model.generate(
        **inputs,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        num_return_sequences=num_return_sequences,
        pad_token_id=tokenizer.eos_token_id # Utilise EOS pour le padding si PAD n'existe pas bien
    )
    
    # 2.3 Décodage et Nettoyage
    output = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return output

# --- 3. EXEMPLE DE TEST ---
# Testons si le modèle peut générer une phrase grammaticalement correcte
TEST_PROMPT = "Quelle est la capitale de la Grèce ? La capitale de la Grèce est"

print(f"\nPrompt de test : '{TEST_PROMPT}'")
print("-" * 50)

generated_text = generate_lamina_text(TEST_PROMPT, max_new_tokens=30)

print(f"\nTexte généré par Lamina Intermédiaire (51M) :")
print(generated_text)

# Deuxième test pour la créativité (comme les "baleines qui inventent les hommes")
CREATIVE_PROMPT = "Les baleines peuvent inventer"

print(f"\nPrompt créatif : '{CREATIVE_PROMPT}'")
print("-" * 50)

creative_text = generate_lamina_text(CREATIVE_PROMPT, max_new_tokens=40)
print(f"\nTexte généré par Lamina Intermédiaire (51M) :")
print(creative_text)