diff --git a/src/main.py b/src/main.py index 01b0969..84f6a07 100644 --- a/src/main.py +++ b/src/main.py @@ -2,6 +2,8 @@ import torch from transformers import AutoModelForCausalLM, AutoTokenizer model_name = "Qwen/Qwen2.5-7B-Instruct" +max_new_tokens = 512 +temperature = 0.1 model = AutoModelForCausalLM.from_pretrained( model_name, @@ -14,14 +16,14 @@ print("Modèle initialisé : "+model_name) print("Pour quitter, entrez \"exit\"") print("") -def generate_response(prompt, max_new_tokens=512): +def generate_response(prompt): chat.append({"role": "user", "content": prompt}) formatted_chat = tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True) inputs = tokenizer(formatted_chat, return_tensors="pt", add_special_tokens=False) inputs = {key: tensor.to(model.device) for key, tensor in inputs.items()} - outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.1) + outputs = model.generate(**inputs, max_new_tokens=max_new_tokens, temperature=temperature) decoded_output = tokenizer.decode(outputs[0][inputs['input_ids'].size(1):], skip_special_tokens=True) chat.append({"role": "assistant", "content": decoded_output}) @@ -36,7 +38,7 @@ formatted_chat = tokenizer.apply_chat_template(chat, tokenize=False, add_generat inputs = tokenizer(formatted_chat, return_tensors="pt", add_special_tokens=False) inputs = {key: tensor.to(model.device) for key, tensor in inputs.items()} -outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.1) +outputs = model.generate(**inputs, max_new_tokens=max_new_tokens, temperature=temperature) decoded_output = tokenizer.decode(outputs[0][inputs['input_ids'].size(1):], skip_special_tokens=True) chat.append({"role": "assistant", "content": decoded_output}) @@ -47,7 +49,8 @@ while True: user_input = input("You: ") if user_input.lower() == "exit": - print("Chatbot: Goodbye!") + print("") + print("Déconnexion") break response = generate_response(user_input)