19 lines
732 B
Python
19 lines
732 B
Python
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))
|