Views
No views yet

| Benchmark | Llama-3.2-3B-Instruct | Llama-3.2-3B-LongCoT |
|---|---|---|
| Math | 35.5 | 52.0 |
| GSM8K | 77.3 | 82.3 |
1import time
2import torch
3from transformers import (
4 AutoModelForCausalLM,
5 AutoTokenizer,
6 TextStreamer,
7)
8
9# Model ID from Hugging Face
10model_id = "Kadins/Llama-3.2-3B-LongCoT"
11
12# Load the pre-trained model with appropriate data type and device mapping
13model = AutoModelForCausalLM.from_pretrained(
14 model_id,
15 torch_dtype=torch.bfloat16, # Use bfloat16 for optimized performance
16 device_map="auto", # Automatically map the model to available devices
17)
18
19# Load the tokenizer associated with the model
20tokenizer = AutoTokenizer.from_pretrained(model_id)
21
22def stream_chat(messages, max_new_tokens=8192, top_p=0.95, temperature=0.6):
23 """
24 Generates a response using streaming inference.
25
26 Args:
27 messages (list): A list of dictionaries containing the conversation prompt.
28 max_new_tokens (int): Maximum number of tokens to generate.
29 top_p (float): Nucleus sampling parameter for controlling diversity.
30 temperature (float): Sampling temperature to control response creativity.
31 """
32 # Prepare the input by applying the chat template and tokenizing
33 inputs = tokenizer.apply_chat_template(
34 messages,
35 tokenize=True,
36 add_generation_prompt=True,
37 return_tensors="pt",
38 return_dict=True, # Ensure the output is a dictionary
39 ).to(model.device) # Move the inputs to the same device as the model
40
41 # Initialize the TextStreamer for real-time output
42 streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
43
44 # Record the start time for performance measurement
45 start_time = time.time()
46
47 # Generate the response using the model's generate method with streaming
48 model.generate(
49 **inputs,
50 max_new_tokens=max_new_tokens,
51 do_sample=True,
52 repetition_penalty=1.1,
53 top_p=top_p,
54 temperature=temperature,
55 streamer=streamer, # Enable streaming of the generated tokens
56 )
57
58 # Calculate and print the total response time
59 total_time = time.time() - start_time
60 print(f"\n--- Response finished in {total_time:.2f} seconds ---")
61
62def chat_loop():
63 """
64 Initiates an interactive chat session with the model.
65 Continuously reads user input and generates model responses until the user exits.
66 """
67 while True:
68 # Initialize the conversation with a system message
69 messages = [
70 {"role": "system", "content": "You are a reasoning expert and helpful assistant."},
71 ]
72
73 # Prompt the user for input
74 user_input = input("\nUser: ")
75 if user_input.strip().lower() in ["exit", "quit"]:
76 print("Exiting chat...")
77 break
78
79 # Append the user's message to the conversation history
80 messages.append({"role": "user", "content": user_input})
81
82 print("Assistant: ", end="", flush=True)
83
84 # Generate and stream the assistant's response
85 stream_chat(messages)
86
87 # Note: Currently, the assistant's reply is streamed directly to the console.
88 # To store the assistant's reply in the conversation history, additional handling is required.
89
90if __name__ == "__main__":
91 # Start the interactive chat loop when the script is executed
92 chat_loop()