Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4model_id = "nik-55/bhagavad-gita-gemma-3-1b-it"
5
6model_kwargs = dict(
7 attn_implementation="eager",
8 device_map="auto",
9)
10
11tokenizer = AutoTokenizer.from_pretrained(model_id)
12model = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)
13
14messages = [
15 {
16 "role": "user",
17 "content": "Who is krishna?",
18 },
19]
20
21inputs = tokenizer.apply_chat_template(
22 messages,
23 add_generation_prompt=True,
24 tokenize=True,
25 return_dict=True,
26 return_tensors="pt",
27).to("cuda")
28
29outputs = model.generate(**inputs, max_new_tokens=500, do_sample=True)
30print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1] :]))
31