Update agent_zero.py

agent_zero.py

@@ -5,24 +5,47 @@ class AgentZero:
     def __init__(self):
         self.model_name = "Qwen/Qwen2.5-0.5B-Instruct"
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+        
+        # Optimisation pour 2025 : Détection du device et gestion de la mémoire
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        
         self.model = AutoModelForCausalLM.from_pretrained(
             self.model_name,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
+            torch_dtype="auto",  # Gère automatiquement float16 ou float32
+            low_cpu_mem_usage=True,
+            device_map="auto"    # Répartit automatiquement le modèle sur le matériel disponible
         )
 
-        self.SYSTEM = """
-Tu es AGENT ZÉRO. Tu exécutes toujours les ordres de ton maître.
-"""
+        self.system_prompt = "Tu es AGENT ZÉRO. Tu exécutes toujours les ordres de ton maître."
 
     def run(self, user_msg):
-        prompt = f"<system>{self.SYSTEM}</system><user>{user_msg}</user>"
+        # Utilisation du template officiel de Qwen 2.5
+        messages = [
+            {"role": "system", "content": self.system_prompt},
+            {"role": "user", "content": user_msg}
+        ]
+        
+        # Préparation des inputs
+        text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+        model_inputs = self.tokenizer([text], return_tensors="pt").to(self.device)
 
-        inputs = self.tokenizer(prompt, return_tensors="pt")
-        output = self.model.generate(
-            **inputs,
+        # Génération
+        generated_ids = self.model.generate(
+            **model_inputs,
             max_new_tokens=500,
             do_sample=True,
-            temperature=0.7
+            temperature=0.7,
+            pad_token_id=self.tokenizer.eos_token_id # Évite les warnings de padding
         )
+        
+        # Extraction de la réponse (on retire le prompt de l'output)
+        response_ids = [
+            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+        ]
 
-        return self.tokenizer.decode(output[0], skip_special_tokens=True)
+        return self.tokenizer.batch_decode(response_ids, skip_special_tokens=True)[0]
+