Here provides a code snippet to show you how to load the tokenizer and model and how to generate contents.
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model_name = "jasonnlp123/TAT-R1"
4
5model = AutoModelForCausalLM.from_pretrained(
6 model_name,
7 torch_dtype="auto",
8 device_map="auto"
9)
10tokenizer = AutoTokenizer.from_pretrained(model_name)
11
12
13system_prompt = """A conversation between User and Assistant. The User asks a question, and the Assistant solves it. \
14The Assistant first thinks about the reasoning process in the mind and then provides the User with the answer. \
15The reasoning process is enclosed within <think> </think> and answer is enclosed within <answer> </answer> tags, respectively, \
16i.e., <think> reasoning process here </think> <answer> answer here </answer>. \
17
18User:
19{}
20
21Assistant:
22"""
23
24# For English to Chinese translation, use:
25query = "Translate the flowing text into Chinese, do not explain：\n{}"
26# For Chinese to English translation, use:
27# query = "Translate the flowing text into English, do not explain：\n{}"
28
29src_text = "Plants make oxygen which humans breathe, and they take in carbon-dioxide which humans exhale (that is, breathe out)."
30prompt = system_prompt.format(query.format(src_text))
31
32
33model_inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
34
35generated_ids = model.generate(
36 **model_inputs,
37 max_new_tokens=2048
38)
39generated_ids = [
40 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
41]
42
43response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
44print(response)
45
