Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3
4model_name = "SuperHotDogCat/jmmlu-gemma-3-12b-thinking"
5tokenizer = AutoTokenizer.from_pretrained(model_name)
6tokenizer.pad_token_id = tokenizer.eos_token_id
7model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto")
8model.eval()
9
10# gemma3p rompt
11prompt_template = '''<start_of_turn>user
12You are a helpful and intelligent AI assistant.
13Please perform reasoning to answer the user's question.
14First, think through the reasoning process step by step, then provide both the reasoning and the answer to the user.
15Enclose the reasoning process and the answer in <think></think> and <answer></answer> tags respectively.
16Please provide your response in the following format:
17<think> Describe your reasoning process here </think>
18<end_of_turn>
19<start_of_turn>user
20{user_input}
21<end_of_turn>
22<start_of_turn>model
23'''
24
25# 推論
26with torch.no_grad():
27 while True:
28 input_text = input("input: ")
29 input_text = prompt_template.format(user_input=input_text)
30 inputs = tokenizer(input_text, return_tensors="pt")
31 outputs = model.generate(
32 input_ids=inputs["input_ids"],
33 attention_mask=inputs["attention_mask"],
34 max_new_tokens=1024,
35 top_k=64,
36 top_p=0.95,
37 )
38 # decode
39 result_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
40 print(result_text)