from transformers import AutoTokenizer, AutoModelForCausalLM

# Nombre del modelo en Hugging Face
model_name = "meta-llama/Llama-3.1"

# Cargar el tokenizer y el modelo
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_8bit=True,  # Para reducir uso de memoria, usa 8-bit (requiere bitsandbytes)
    device_map="auto"   # Distribuye automáticamente en los GPUs disponibles
)

# Probar el modelo con una entrada
prompt = "¿Qué es Llama 3.1?"
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")  # Asegúrate de que se ejecute en GPU
outputs = model.generate(**inputs, max_length=50)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))