Add variables

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)