Update app.py

app.py

@@ -1,18 +1,31 @@
 import os
 from dotenv import load_dotenv
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
-# Carrega variáveis de ambiente (para testes locais; no Space, usa secrets)
+# Carrega variáveis de ambiente
 load_dotenv()
 api_key = os.getenv("HF_API_TOKEN")
-
-# Verifica se o token existe
 if not api_key:
-    raise ValueError("Erro: HF_API_TOKEN não encontrado. Configure-o nas secrets do Space (Settings > Repository secrets).")
+    raise ValueError("Erro: HF_API_TOKEN não encontrado. Configure-o nas secrets do Space.")
+
+# Configuração de quantização
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype="bfloat16"
+)
 
-# Inicializa o cliente de inferência com o modelo Mistral
-client = InferenceClient(token=api_key, model="mistralai/Mistral-7B-Instruct-v0.3")
+# Inicializa o modelo e o tokenizer
+model_id = "Qwen/Qwen2-7B-Instruct-GPTQ"
+tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=api_key)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    quantization_config=quant_config,
+    device_map="auto",
+    use_auth_token=api_key
+)
 
 # Função do chatbot
 def chat_with_llm(message, history):
@@ -25,22 +38,27 @@ def chat_with_llm(message, history):
                 messages.append({"role": "assistant", "content": bot_msg})
         messages.append({"role": "user", "content": message})
 
-        # Chama a API do Mistral
-        response = client.chat.completions.create(
-            model="mistralai/Mistral-7B-Instruct-v0.3",
-            messages=messages,
-            max_tokens=150,
-            temperature=0.7,
+        # Tokeniza a entrada
+        inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to("cuda")
+        
+        # Gera a resposta
+        outputs = model.generate(
+            inputs,
+            max_new_tokens=500,
+            temperature=0.5,
+            top_p=0.9,
+            do_sample=True
         )
-        return response.choices[0].message.content
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response
     except Exception as e:
-        return f"Erro: {str(e)}. Verifique o token API, conexão com a internet ou limites da API gratuita."
+        return f"Erro: {str(e)}. Verifique o token API, conexão com a internet ou requisitos de hardware."
 
-# Cria a interface do Gradio
+# Interface Gradio
 demo = gr.ChatInterface(
     fn=chat_with_llm,
-    title="Chatbot com Mistral (Gratuito via HF API)",
-    description="Teste o chatbot usando Mistral. Plano gratuito com limites.",
+    title="Chatbot com Qwen2-7B Quantizado",
+    description="Chatbot usando Qwen2-7B quantizado em 4-bit para maior eficiência.",
 )
 
 if __name__ == "__main__":