Update app.py

app.py

@@ -1,25 +1,28 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import threading
+import torch
 
-# Inicializar o modelo e o tokenizer
-model_name = "lambdaindie/lambda-1v-1B"  # Troca com o nome do modelo que estás a usar
-model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")  # Usando GPU (ou "cpu" se não tiveres GPU)
+# Detectar dispositivo automaticamente (GPU ou CPU)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Carregar modelo e tokenizer
+model_name = "lambdaindie/lambda-1v-1B"
+model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 stop_flag = {"stop": False}
 
+# Função de resposta
 def respond(prompt, history):
     stop_flag["stop"] = False
 
-    # Prompt modificado conforme solicitado
     full_prompt = f"\nThink a bit step-by-step before answering.  \nQuestion: {prompt} \nAnswer:"
-
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
 
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
-    # Thread para geração de texto
+    # Iniciar thread de geração
     generation_thread = threading.Thread(
         target=model.generate,
         kwargs={
@@ -40,27 +43,29 @@ def respond(prompt, history):
         if stop_flag["stop"]:
             return "", history
         reasoning += new_text
-        yield "", history[:-1] + [(prompt, f"<div class='final-answer'>{reasoning}</div>")]
+        yield "", (history or []) + [(prompt, f"<div class='final-answer'>{reasoning}</div>")]
 
+# Função para parar a geração
 def stop_generation():
     stop_flag["stop"] = True
 
-# Definir a interface do Gradio
+# Interface Gradio
 with gr.Blocks(css="""
-    #chatbot, .gr-markdown, .gr-button, .gr-textbox {  
-        font-family: 'JetBrains Mono', monospace !important;  
-        font-size: 11px !important;  
-    }  
-    .final-answer {  
-        background-color: #1e1e1e;  
-        color: #ffffff;  
-        padding: 10px;  
-        border-left: 4px solid #4caf50;  
-        font-family: 'JetBrains Mono', monospace !important;  
-        white-space: pre-wrap;  
-        font-size: 11px !important;  
+    #chatbot, .gr-markdown, .gr-button, .gr-textbox {
+        font-family: 'JetBrains Mono', monospace !important;
+        font-size: 11px !important;
+    }
+    .final-answer {
+        background-color: #1e1e1e;
+        color: #ffffff;
+        padding: 10px;
+        border-left: 4px solid #4caf50;
+        font-family: 'JetBrains Mono', monospace !important;
+        white-space: pre-wrap;
+        font-size: 11px !important;
     }
 """) as demo:
+    gr.Markdown('<link href="https://fonts.googleapis.com/css2?family=JetBrains+Mono&display=swap" rel="stylesheet">')
     gr.Markdown("## λambdAI — Reasoning Chat")
 
     chatbot = gr.Chatbot(elem_id="chatbot")