matsuo-llm-course
Collection
16 items
•
Updated
このモデルは東京大学松尾・岩澤研究室のLLM講座2024の課題のために作られたものです。
model_name = "miya-99999/matsuo_llm_exp002"
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
model = AutoModelForCausalLM.from_pretrained(
model_name,
device_map=device,
token=HF_TOKEN,
torch_dtype=torch.bfloat16,
use_cache=True
)
chat = [
{"role": "user", "content": "こんにちは。いい天気ですね。"},
]
tokenized_input = tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=True, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(
tokenized_input,
max_new_tokens=512,
do_sample=False,
repetition_penalty=1.2,
pad_token_id=tokenizer.pad_token_id,
eos_token_id=[1,107],
)[0]
output = tokenizer.decode(outputs[tokenized_input.size(1):], skip_special_tokens=True)