Update handler.py

handler.py

@@ -1,30 +1,33 @@
 from typing import Dict, Any
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import LlamaForCausalLM, LlamaTokenizer
 import torch
 
 class EndpointHandler:
     def __init__(self, path=""):
         # Configuração do modelo
-        self.model_name_or_path = "souzat19/Llama3.1_fn14133.29122024"
+        self.model_name_or_path = path or "souzat19/Llama3.1_fn14133.29122024"
         
         # Detecta se GPU está disponível
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         print(f"Using device: {self.device}")
         
         print("Initializing tokenizer...")
-        # Inicialização do tokenizer
-        self.tokenizer = AutoTokenizer.from_pretrained(
+        self.tokenizer = LlamaTokenizer.from_pretrained(
             self.model_name_or_path,
             trust_remote_code=True
         )
         
         print("Initializing model...")
-        # Inicialização do modelo com configurações mínimas
-        self.model = AutoModelForCausalLM.from_pretrained(
+        self.model = LlamaForCausalLM.from_pretrained(
             self.model_name_or_path,
             torch_dtype=torch.float32,
-            trust_remote_code=True
-        ).to(self.device)  # Move para GPU se disponível
+            trust_remote_code=True,
+            device_map="auto" if torch.cuda.is_available() else None,
+            local_files_only=True if path else False
+        )
+        
+        if not torch.cuda.is_available():
+            self.model = self.model.to("cpu")
         
         print("Model initialized successfully")
         
@@ -52,11 +55,12 @@ Você é um assistente especializado em planejamento de compras públicas de aco
                 formatted_prompt,
                 return_tensors="pt",
                 truncation=True,
-                max_length=4096
+                max_length=4096,
+                add_special_tokens=True
             )
             
-            # Move input para mesmo device do modelo
-            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            if torch.cuda.is_available():
+                inputs = {k: v.to(self.device) for k, v in inputs.items()}
 
             # Gera a resposta
             with torch.no_grad():
@@ -66,7 +70,9 @@ Você é um assistente especializado em planejamento de compras públicas de aco
                     temperature=0.5,
                     top_p=0.95,
                     top_k=50,
-                    do_sample=True
+                    do_sample=True,
+                    pad_token_id=self.tokenizer.pad_token_id,
+                    eos_token_id=self.tokenizer.eos_token_id
                 )
 
             # Decodifica a resposta