Update handler.py

handler.py

@@ -11,8 +11,9 @@ class EndpointHandler:
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path)
         self.model = AutoModelForCausalLM.from_pretrained(
             self.model_name_or_path,
-            torch_dtype=torch.float16,
-            device_map="auto"
+            device_map="auto",
+            trust_remote_code=True,
+            use_cache=True
         )
         
         # Template do prompt no formato Alpaca
@@ -25,16 +26,6 @@ Você é um assistente especializado em planejamento de compras públicas de aco
 """
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Processa a entrada e retorna a resposta do modelo.
-        
-        Args:
-            data: Dicionário contendo a entrada do usuário
-                 Formato esperado: {"text": "pergunta do usuário"}
-        
-        Returns:
-            Dict contendo a resposta do modelo
-        """
         try:
             # Extrai o texto da entrada
             input_text = data.get("text", "")
@@ -52,16 +43,17 @@ Você é um assistente especializado em planejamento de compras públicas de aco
             inputs = inputs.to(self.model.device)
 
             # Gera a resposta
-            outputs = self.model.generate(
-                **inputs,
-                max_new_tokens=2096,
-                temperature=0.5,
-                top_p=0.95,
-                top_k=50,
-                pad_token_id=self.tokenizer.pad_token_id,
-                eos_token_id=self.tokenizer.eos_token_id,
-                do_sample=True
-            )
+            with torch.no_grad():
+                outputs = self.model.generate(
+                    **inputs,
+                    max_new_tokens=2096,
+                    temperature=0.5,
+                    top_p=0.95,
+                    top_k=50,
+                    do_sample=True,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id
+                )
 
             # Decodifica a resposta
             response_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
@@ -91,5 +83,4 @@ Você é um assistente especializado em planejamento de compras públicas de aco
         """
         if not text or len(text.strip()) == 0:
             return False
-        return True
-        
+        return True