Update app.py

app.py

@@ -77,6 +77,82 @@ def get_llm():
     return HuggingFacePipeline(pipeline=text_pipeline)
 
 
+def get_llm():
+    """Carrega o LLM quantizado localmente usando bitsandbytes e um pipeline com chat template."""
+    # Atualizando o nome do modelo para carregar o tokenizer correto,
+    # mesmo que o arquivo GGUF seja de outro repo.
+    # O tokenizer ainda deve ser compatível com o modelo base "google/gemma-2b-it".
+    TOKENIZER_MODEL = "google/gemma-2b-it"
+    model_path = os.path.join(CACHE_FOLDER, LOCAL_MODEL_FILE)
+    logger.info(f"Carregando LLM quantizado localmente: {LOCAL_MODEL_FILE}")
+
+    try:
+        # Configuração da quantização com bitsandbytes
+        # `load_in_4bit=True` habilita a quantização de 4 bits
+        # `bnb_4bit_quant_type` define o tipo de quantização (fp4 ou nf4)
+        # `bnb_4bit_compute_dtype` define o tipo de dados para o cálculo (float16 é comum para GPUs)
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            llm_int8_threshold=6.0, # Adicionado para compatibilidade
+            llm_int8_skip_modules=None, # Adicionado para compatibilidade
+            llm_int8_enable_fp32_cpu_offload=False # Adicionado para compatibilidade
+        )
+
+        # Carregando o tokenizer e o modelo, aplicando a quantização
+        tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_MODEL, cache_dir=CACHE_FOLDER)
+        # O modelo Gemma-2b-it possui um template de chat embutido
+        if tokenizer.chat_template is None:
+            logger.warning("O modelo não tem um template de chat. Usando o template padrão.")
+            tokenizer.chat_template = "{% for message in messages %}{% if message['role'] == 'user' %}{{ '[User]: ' + message['content'] + '\n\n' }}{% else %}{{ '[Assistant]: ' + message['content'] + '\n\n' }}{% endif %}{% endfor %}"
+
+
+        # O modelo a ser carregado agora é o arquivo GGUF local.
+        # No entanto, bitsandbytes é geralmente usado para carregar modelos transformers
+        # não GGUF. Para GGUF, você geralmente usaria uma biblioteca como `ctransformers`
+        # ou `llama-cpp-python`.
+        # Como o código original usava bitsandbytes com um nome de arquivo GGUF,
+        # assumirei que a intenção era carregar um modelo compatível com transformers/bitsandbytes,
+        # talvez com um nome de arquivo .safetensors ou .bin.
+        # Vou reverter para carregar o modelo diretamente do repo original
+        # "google/gemma-2b-it" usando bitsandbytes, já que o arquivo GGUF
+        # não é o formato esperado para bitsandbytes/transformers.
+        # Se a intenção REALMENTE for usar o arquivo GGUF, a abordagem de carregamento
+        # precisará ser completamente reescrita usando uma biblioteca apropriada (ex: ctransformers).
+
+        logger.warning("Detectado uso de bitsandbytes com nome de arquivo .gguf. Bitsandbytes é para modelos transformers (ex: .safetensors, .bin). Revertendo para carregar o modelo diretamente do repo original 'google/gemma-2b-it' com bitsandbytes.")
+
+        model = AutoModelForCausalLM.from_pretrained(
+            TOKENIZER_MODEL, # Carregando do repo original para usar bitsandbytes
+            cache_dir=CACHE_FOLDER,
+            device_map="auto", # Tenta usar a GPU, se disponível
+            quantization_config=quantization_config # Adicionando a configuração de quantização
+        )
+
+        # Criando a pipeline de geração de texto
+        text_pipeline = pipeline(
+            "text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_new_tokens=1024,
+            temperature=0.3,
+            # Parâmetros de parada para evitar que o modelo continue a gerar após a resposta
+            # Note: stop_sequence pode não ser universalmente suportado por todas as pipelines/modelos
+            # dependendo da implementação específica.
+            eos_token_id=tokenizer.eos_token_id # Usar EOS token é mais robusto
+        )
+
+        # Retornando a LLM da LangChain que usa a pipeline
+        return HuggingFacePipeline(pipeline=text_pipeline)
+
+    except Exception as e:
+        logger.error(f"Erro ao carregar o modelo. Erro: {e}")
+        # Informar o usuário sobre a incompatibilidade potencial
+        if "bitsandbytes" in str(e).lower() and ".gguf" in LOCAL_MODEL_FILE.lower():
+             logger.error("Possível erro de incompatibilidade: bitsandbytes é usado para modelos transformers, não GGUF. Considere usar uma biblioteca como 'ctransformers' ou 'llama-cpp-python' para arquivos GGUF.")
+        raise e
+
 
 def get_embedding_model():
     return HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL, cache_folder=CACHE_FOLDER)