Spaces:
Running
Running
from transformers import AutoModelForCausalLM, AutoTokenizer | |
# Load model and tokenizer | |
model_name = "prithivMLmods/QwQ-LCoT-7B-Instruct" | |
model = AutoModelForCausalLM.from_pretrained( | |
model_name, | |
torch_dtype="auto", | |
device_map="auto" | |
) | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
# Initialize persistent conversation with a system message | |
system_message = {"role": "system", "content": "You are a helpful and harmless assistant. You are Qwen developed by Alibaba. You should think step-by-step."} | |
messages = [system_message] | |
# Chat loop to maintain persistence | |
while True: | |
user_input = input("User: ") # Get user input | |
if user_input.lower() in {"exit", "quit"}: | |
print("Chat session ended.") | |
break | |
# Append user message to the conversation history | |
messages.append({"role": "user", "content": user_input}) | |
# Format the messages for the model | |
text = tokenizer.apply_chat_template( | |
messages, | |
tokenize=False, | |
add_generation_prompt=True | |
) | |
model_inputs = tokenizer([text], return_tensors="pt").to(model.device) | |
# Generate response | |
generated_ids = model.generate( | |
**model_inputs, | |
max_new_tokens=512 | |
) | |
generated_ids = [ | |
output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids) | |
] | |
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] | |
# Append assistant's response to the conversation history | |
messages.append({"role": "assistant", "content": response}) | |
# Display the assistant's response | |
print(f"Assistant: {response}") | |