diff --git a/src/main.py b/src/main.py index 84f6a07..623a2f8 100644 --- a/src/main.py +++ b/src/main.py @@ -2,7 +2,7 @@ import torch from transformers import AutoModelForCausalLM, AutoTokenizer model_name = "Qwen/Qwen2.5-7B-Instruct" -max_new_tokens = 512 +max_new_tokens = 2048 temperature = 0.1 model = AutoModelForCausalLM.from_pretrained( @@ -43,10 +43,10 @@ decoded_output = tokenizer.decode(outputs[0][inputs['input_ids'].size(1):], skip chat.append({"role": "assistant", "content": decoded_output}) -print(decoded_output) +print("Assistant: ", decoded_output) while True: - user_input = input("You: ") + user_input = input("User: ") if user_input.lower() == "exit": print("") @@ -54,5 +54,5 @@ while True: break response = generate_response(user_input) - print("Chatbot:", response) + print("Assistant: ", response)