Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3model_id = "danialbha7899/mistral-7b-shakespeare-chat"
4
5tokenizer = AutoTokenizer.from_pretrained(model_id)
6model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
7
8messages = [
9 {"role": "user", "content": "How is the weather today?"}
10]
11
12inputs = tokenizer.apply_chat_template(
13 messages,
14 return_tensors="pt",
15 add_generation_prompt=True
16).to("cuda")
17
18outputs = model.generate(**inputs, max_new_tokens=150)
19
20print(tokenizer.decode(outputs[0], skip_special_tokens=True))