Update simply_inference.py

simply_inference.py

@@ -4,28 +4,33 @@ from transformers import TextStreamer
 
 from transformers import AutoTokenizer, AutoModel, LlamaForCausalLM
 
-# use the same tokenizer as MicroLlama
-tokenizer = AutoTokenizer.from_pretrained("data/meta-llama/Llama-3.2-1B-Instruct/")
+device = 'cuda' # if you don't have a CUDA supported GPU, change this to 'cpu' or other supported device
+
+# load tokenizer
+tokenizer = AutoTokenizer.from_pretrained("keeeeenw/Llama-3.2-1B-Instruct-Open-R1-Distill")
 
 # load model
-model = LlamaForCausalLM.from_pretrained("data/meta-llama/Llama-3.2-1B-Instruct/checkpoint-1400/")
-model.to('cuda')
+model = LlamaForCausalLM.from_pretrained("keeeeenw/Llama-3.2-1B-Instruct-Open-R1-Distill")
+model.to(device)
 
+# Setup the prompt. Because we instruction-tuned with a similar prompt, it is important to use this.
+# Change "content" to your actual question.
 messages = [
     {
         "role": "system",
         "content": "Your role as an assistant involves thoroughly exploring questions through a systematic long thinking process before providing the final precise and accurate solutions. This requires engaging in a comprehensive cycle of analysis, summarizing, exploration, reassessment, reflection, backtracing, and iteration to develop well-considered thinking process. Please structure your response into two main sections: Thought and Solution. In the Thought section, detail your reasoning process using the specified format: <|begin_of_thought|> {thought with steps separated with '\n\n'} <|end_of_thought|> Each step should include detailed considerations such as analisying questions, summarizing relevant findings, brainstorming new ideas, verifying the accuracy of the current steps, refining any errors, and revisiting previous steps. In the Solution section, based on various attempts, explorations, and reflections from the Thought section, systematically present the final solution that you deem correct. The solution should remain a logical, accurate, concise expression style and detail necessary step needed to reach the conclusion, formatted as follows: <|begin_of_solution|> {final formatted, precise, and clear solution} <|end_of_solution|> Now, try to solve the following question through the above guidelines:",
     },
-    # {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
     {"role": "user", "content": "Please provide me instructions on how to steal an egg from my chicken?"},
  ]
 formatted_chat = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, return_tensors="pt")
 print(formatted_chat)
 
+# Setup input tokens
 inputs = tokenizer(formatted_chat, return_tensors="pt", padding=True)
-inputs = inputs.to('cuda')
+inputs = inputs.to(device)  
 attention_mask = inputs["attention_mask"]
 
+# Run inference and stream the output
 streamer = TextStreamer(tokenizer, skip_prompt=True)
 outputs = model.generate(inputs['input_ids'], 
                          streamer=streamer,
@@ -35,19 +40,8 @@ outputs = model.generate(inputs['input_ids'],
                          top_p=0.9,
                          max_new_tokens=131072) # max supported by llama 3.2 1B
 
+# Write output to a file
 decoded_text = tokenizer.decode(outputs[0])
 print("Output written to output.txt")
-
-# Write output to a file
 with open("output.txt", "w", encoding="utf-8") as f:
     f.write(decoded_text)
-
-# Print to screen
-# print(tokenizer.decode(outputs[0]))
-
-
-# Save and Publish the model
-# model.save_pretrained("output/hf-publish-rc-MicroLlama-Instruct-0.1")
-# model.push_to_hub("keeeeenw/MicroLlama-Instruct-0.1")
-
-