If You had an Error trying to run with Transformers use this
#1
by jsob7 - opened
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
MODEL = "nvidia/Riva-Translate-4B-Instruct-v2"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
model = AutoModelForCausalLM.from_pretrained(
MODEL,
torch_dtype=torch.float16,
device_map="auto",
)
messages = [
{
"role": "system",
"content": "[tag-pair]",
},
{
"role": "user",
"content": "The GRACE mission is a collaboration between the NASA and German Aerospace Center.?",
},
]
inputs = tokenizer.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt",
).to(model.device)
outputs = model.generate(
**inputs,
max_new_tokens=128,
pad_token_id=tokenizer.eos_token_id,
)
# Remove os tokens do prompt original
input_length = inputs["input_ids"].shape[-1]
response = tokenizer.decode(
outputs[0][input_length:],
skip_special_tokens=True,
)
print(response)
jsob7 changed discussion title from If Your had an Error trying to run with Transformers use this to If You had an Error trying to run with Transformers use this