Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model_name = "AdvRahul/Axion-Pro-Indic-24B"
4
5tokenizer = AutoTokenizer.from_pretrained(model_name)
6model = AutoModelForCausalLM.from_pretrained(
7 model_name, torch_dtype="auto", device_map="auto"
8)
9
10prompt = "Who are you and what is your purpose on this planet?"
11
12messages = [{"role": "user", "content": prompt}]
13text = tokenizer.apply_chat_template(
14 messages,
15 tokenize=False,
16 enable_thinking=True, # Default True; set False for no-think mode
17)
18
19model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
20
21generated_ids = model.generate(**model_inputs, max_new_tokens=8192)
22output_ids = generated_ids[0][len(model_inputs.input_ids[0]) :].tolist()
23output_text = tokenizer.decode(output_ids)
24
25if "</think>" in output_text:
26 reasoning_content = output_text.split("</think>")[0].rstrip("\n")
27 content = output_text.split("</think>")[-1].lstrip("\n").rstrip("</s>")
28else:
29 reasoning_content = ""
30 content = output_text.rstrip("</s>")
31
32print("reasoning content:", reasoning_content)
33print("content:", content)