Update modules/local_llm.py

modules/local_llm.py

@@ -1,108 +1,181 @@
-from llama_cpp import Llama
 import os
-from loguru import logger
 import threading
+from loguru import logger
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+# Removida a dependência 'llama_cpp'
 
-# CAMINHOS NO HF SPACES (CORRIGIDOS PARA O AMBIENTE ATUAL)
-FINETUNED_PATH = "/home/user/data/finetuned_phi3"        # LoRA angolano
-# CORREÇÃO: O Dockerfile baixa para /home/user/models/
+# === Variáveis de Ambiente e Caminhos (Mantidas, mas adaptadas) ===
+# NOTA IMPORTANTE: Para usar 'transformers', o modelo será carregado do Hugging Face.
+# Os caminhos GGUF e FINETUNED_PATH serão usados APENAS para verificação/logging,
+# mas o modelo carregado será o nativo do HF.
+FINETUNED_PATH = "/home/user/data/finetuned_phi3"
 GGUF_FILENAME = "Phi-3-mini-4k-instruct.Q4_K_M.gguf"
 GGUF_PATH = f"/home/user/models/{GGUF_FILENAME}"
+HF_MODEL_ID = "microsoft/Phi-3-mini-4k-instruct"
+# =================================================================
 
 class Phi3LLM:
+    # Usaremos uma tupla (model, tokenizer)
     _llm = None
     _available_checked = False
-    _is_available = False
-    MODEL_ID = "PHI-3 3.8B"
-    MODEL_SIZE_RAM_GB = "~7-8GB"
-
+    _is_available = False # True se as bibliotecas estiverem OK
+    MODEL_ID = "PHI-3 3.8B (HF Transformers)"
+    MODEL_SIZE_RAM_GB = "~7-8GB (4-bit: ~4GB)"
+    
+    # 1. Checagem de disponibilidade (Adaptada)
     @classmethod
     def is_available(cls) -> bool:
         """
-        VERIFICA SE O MODELO GGUF EXISTE → CACHEIA O RESULTADO
-        EVITA CHAMADAS REPETIDAS AO DISCO → MAIS RÁPIDO NO START
+        VERIFICA SE O AMBIENTE ESTÁ PRONTO PARA CARREGAR O MODELO HF.
         """
         if not cls._available_checked:
-            cls._is_available = os.path.isfile(GGUF_PATH)
-            cls._available_checked = True
-            if cls._is_available:
-                logger.info(f"{cls.MODEL_ID} GGUF ENCONTRADO → {GGUF_PATH}")
-            else:
-                logger.warning(f"{cls.MODEL_ID} GGUF NÃO ENCONTRADO! Caminho: {GGUF_PATH}")
-                logger.warning("AKIRA VAI USAR MISTRAL/GEMINI COMO FALLBACK")
+            # Não verifica mais o arquivo GGUF, apenas a capacidade de execução
+            try:
+                import torch
+                from transformers import AutoModelForCausalLM, AutoTokenizer
+                cls._is_available = True
+                cls._available_checked = True
+                logger.info(f"{cls.MODEL_ID} AMBIENTE DE EXECUÇÃO PRONTO (PyTorch/Transformers).")
+                
+                if os.path.isfile(GGUF_PATH):
+                    logger.warning("GGUF ENCONTRADO, MAS SERÁ IGNORADO. O modelo será carregado do HF para compatibilidade com 'transformers'.")
+                else:
+                    logger.warning(f"GGUF NÃO ENCONTRADO EM: {GGUF_PATH}")
+                    
+            except ImportError as e:
+                cls._is_available = False
+                cls._available_checked = True
+                logger.error(f"FALHA DE DEPENDÊNCIA: {e}. Certifique-se de que 'torch', 'transformers' e 'accelerate' estão instalados.")
         return cls._is_available
 
+    # 2. Carregamento do Modelo (Adaptado para Hugging Face)
     @classmethod
     def _get_llm(cls):
         """
-        CARREGA O MODELO UMA ÚNICA VEZ → SINGLETON + mlock
-        PHI-3 é menor (3.8B) → Deve usar menos RAM e ser mais rápido.
+        CARREGA O MODELO (MODELO + TOKENIZER) UMA ÚNICA VEZ.
+        Utiliza quantização 4-bit se GPU estiver disponível.
         """
         if cls._llm is None and cls.is_available():
+            
+            # Checa por GPU
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+            print(f"Dispositivo de inferência selecionado: {device.upper()}")
+            
             try:
-                logger.info(f"CARREGANDO {cls.MODEL_ID} GGUF Q4_K_M → RAM: {cls.MODEL_SIZE_RAM_GB}, CPU MÍNIMA!")
-                cls._llm = Llama(
-                    model_path=GGUF_PATH,
-                    n_ctx=4096,         # Contexto gigante
-                    n_batch=512,        # Batch grande = menos CPU
-                    n_threads=2,        # Só 2 threads → CPU em paz!
-                    n_gpu_layers=0,     # Tudo na RAM
-                    use_mlock=True,     # Trava na RAM física → nunca swap!
-                    verbose=False,
-                    n_parts=1,          # Carrega tudo de uma vez
-                    seed=-1,
-                    logits_all=True,
+                logger.info(f"CARREGANDO {cls.MODEL_ID} DO HF → DEVICE: {device.upper()}")
+
+                # Configuração de Quantização (equivalente a otimização de RAM)
+                if device == "cuda":
+                    logger.info("Configuração de quantização 4-bit (BitsAndBytes) ativada para economia de VRAM.")
+                    bnb_config = BitsAndBytesConfig(
+                        load_in_4bit=True,
+                        bnb_4bit_quant_type="nf4",
+                        bnb_4bit_compute_dtype=torch.bfloat16,
+                    )
+                else:
+                    bnb_config = None
+                    logger.info("Rodando em CPU. Sem quantização 4-bit (usará float32/float16).")
+                
+                # 1. Carrega o Tokenizer
+                tokenizer = AutoTokenizer.from_pretrained(HF_MODEL_ID, trust_remote_code=True)
+                
+                # 2. Carrega o Modelo
+                model = AutoModelForCausalLM.from_pretrained(
+                    HF_MODEL_ID,
+                    torch_dtype=torch.bfloat16 if device == "cuda" else torch.float32,
+                    trust_remote_code=True,
+                    quantization_config=bnb_config,
+                    device_map="auto" # Mapeamento automático (substitui n_gpu_layers/mlock)
                 )
-                logger.success(f"{cls.MODEL_ID} GGUF CARREGADO COM SUCESSO → {cls.MODEL_SIZE_RAM_GB} RAM USADA!")
 
-                # =================================================================
-                # TENTA CARREGAR LORA SE EXISTIR (funcionalidade depende da versão llama-cpp-python)
-                # =================================================================
+                # LoRA loading is complex and dependent on HF Trainer/Peft setup.
+                # Here, we only log the path, as dynamic LoRA loading in 'transformers'
+                # is not as straightforward as in llama-cpp.
                 lora_path = os.path.join(FINETUNED_PATH, "lora_leve")
                 if os.path.isdir(lora_path):
-                     logger.warning(f"LoRA encontrado em {lora_path}, mas 'load_lora' não é suportado pela versão do llama-cpp-python. A carregar modelo base.")
+                    logger.warning(f"LoRA encontrado em {lora_path}, mas o carregamento não é suportado pelo método de inferência direta com AutoModel. Usando modelo base.")
                 else:
-                     logger.info("Usando modelo GGUF base (sem LoRA).")
+                    logger.info("Usando modelo base (sem LoRA).")
 
+                cls._llm = (model, tokenizer)
+                logger.success(f"{cls.MODEL_ID} CARREGADO COM SUCESSO. Config: {device.upper()} | 4-bit: {bnb_config is not None}")
+            
             except Exception as e:
-                logger.error(f"ERRO CRÍTICO AO CARREGAR {cls.MODEL_ID} GGUF: {e}")
+                logger.error(f"ERRO CRÍTICO AO CARREGAR {cls.MODEL_ID} DO HF: {e}")
                 import traceback
                 logger.error(traceback.format_exc())
-                cls._llm = None # Garante que não tente de novo
+                cls._llm = None
+                
         return cls._llm
 
+    # 3. Geração de Resposta (Adaptado)
     @classmethod
     def generate(cls, prompt: str, max_tokens: int = 60) -> str:
         """
-        GERA RESPOSTA COM PHI-3 LOCAL
-        max_tokens=60 → Deve ser mais rápido que o Hermes (7B)
+        GERA RESPOSTA COM PHI-3 USANDO TRANSFORMERS.
+        max_tokens → 'max_new_tokens' na função generate.
         """
-        llm = cls._get_llm()
-        if llm is None:
-            raise RuntimeError(f"{cls.MODEL_ID} não está disponível ou falhou ao carregar")
-
+        llm_pair = cls._get_llm()
+        if llm_pair is None:
+            raise RuntimeError(f"{cls.MODEL_ID} não está disponível ou falhou ao carregar.")
+        
+        model, tokenizer = llm_pair
+        device = model.device # Pega o dispositivo onde o modelo está
+        
         try:
             # FORMATO DE CHAT PHI-3: <|user|>PROMPT<|end|><|assistant|>
-            formatted_prompt = f"<|user|>\n{prompt}<|end|>\n<|assistant|>"
+            # O tokenizer.apply_chat_template faz isso de forma canônica
+            messages = [{"role": "user", "content": prompt}]
+            formatted_prompt = tokenizer.apply_chat_template(
+                messages, 
+                tokenize=False, 
+                add_generation_prompt=True
+            )
 
+            # Tokeniza a entrada
+            input_ids = tokenizer.encode(formatted_prompt, return_tensors="pt").to(device)
+            
             logger.info(f"[{cls.MODEL_ID} LOCAL] Gerando resposta → {max_tokens} tokens")
-            output = llm(
-                formatted_prompt,
-                max_tokens=max_tokens,
-                temperature=0.72,
-                top_p=0.92,
-                stop=["<|end|>", "<|assistant|>", "<|user|>", "\n\n"],
-                echo=False,
-            )
-            text = output["choices"][0]["text"].strip()
             
-            # Limpeza adicional para garantir que não haja tags residuais
-            if text.startswith("<|assistant|>"):
-                text = text[len("<|assistant|>"):].lstrip()
+            with torch.no_grad():
+                output = model.generate(
+                    input_ids,
+                    max_new_tokens=max_tokens,
+                    temperature=0.72,
+                    top_p=0.92,
+                    do_sample=True, # Adicionado para usar temperature/top_p
+                    pad_token_id=tokenizer.eos_token_id, # Boa prática
+                )
+            
+            # Decodifica, ignorando o prompt de entrada
+            text = tokenizer.decode(output[0][input_ids.shape[-1]:], skip_special_tokens=True).strip()
+            
+            # Limpeza manual de tags Phi-3 (apesar do skip_special_tokens, é bom garantir)
+            for tag in ["<|end|>", "<|assistant|>", "<|user|>"]:
+                 text = text.replace(tag, "").strip()
             
             logger.success(f"{cls.MODEL_ID} RESPONDEU → {len(text)} chars")
             return text
 
         except Exception as e:
-            logger.error(f"ERRO NA GERAÇÃO COM {cls.MODEL_ID}: {e}")
-            raise
+            logger.error(f"ERRO NA GERAÇÃO COM {cls.MODEL_ID} (Transformers): {e}")
+            raise
+
+# === Teste Básico de Disponibilidade no Script ===
+if __name__ == '__main__':
+    if Phi3LLM.is_available():
+        print("\nTeste de Carga do Modelo (Pode demorar)...")
+        if Phi3LLM._get_llm():
+            try:
+                # Exemplo de uso
+                prompt = "Me diga três fatos interessantes sobre a capital de Angola."
+                print(f"\nPROMPT: {prompt}")
+                response = Phi3LLM.generate(prompt, max_tokens=100)
+                print(f"\nRESPOSTA DO LLM:\n{response}")
+            except RuntimeError as e:
+                print(f"Erro durante a geração: {e}")
+        else:
+            print("Falha ao carregar o modelo.")
+    else:
+        print("Ambiente não está pronto. Verifique as dependências.")