Update app.py

app.py

@@ -1,36 +1,76 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from peft import PeftModel, PeftConfig
-import torch
-
-# Carrega config do adaptador
-config = PeftConfig.from_pretrained("lambdaindie/lambda-1v-1B")
-
-# Carrega modelo base + LoRA
-base_model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)
-model = PeftModel.from_pretrained(base_model, "lambdaindie/lambda-1v-1B")
-
-# Tokenizer
-tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
-
-# Envia pra CPU
-device = torch.device("cpu")
-model.to(device)
-
-def responder(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=50,
-        temperature=0.8,
-        top_p=0.95,
-        pad_token_id=tokenizer.eos_token_id,
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import threading
+
+# 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)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+stop_flag = {"stop": False}
+
+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)
+
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    # Thread para geração de texto
+    generation_thread = threading.Thread(
+        target=model.generate,
+        kwargs={
+            "input_ids": inputs["input_ids"],
+            "attention_mask": inputs["attention_mask"],
+            "max_new_tokens": 512,
+            "do_sample": True,
+            "temperature": 0.7,
+            "top_p": 0.9,
+            "pad_token_id": tokenizer.eos_token_id,
+            "streamer": streamer,
+        }
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-iface = gr.Interface(fn=responder,
-                     inputs=gr.Textbox(lines=2, placeholder="Escreve algo..."),
-                     outputs="text",
-                     title="Lambda-1v-1B (LoRA)",
-                     description="Modelo LoRA fine-tuned por Marius Jabami.")
-iface.launch()
+    generation_thread.start()
+
+    reasoning = ""
+    for new_text in streamer:
+        if stop_flag["stop"]:
+            return "", history
+        reasoning += new_text
+        yield "", history[:-1] + [(prompt, f"<div class='final-answer'>{reasoning}</div>")]
+
+def stop_generation():
+    stop_flag["stop"] = True
+
+# Definir a interface do 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;  
+    }
+""") as demo:
+    gr.Markdown("## λambdAI — Reasoning Chat")
+
+    chatbot = gr.Chatbot(elem_id="chatbot")
+    with gr.Row():
+        txt = gr.Textbox(placeholder="Digite sua pergunta...", show_label=False)
+        send_btn = gr.Button("Enviar")
+        stop_btn = gr.Button("Parar")
+
+    send_btn.click(respond, [txt, chatbot], [txt, chatbot])
+    txt.submit(respond, [txt, chatbot], [txt, chatbot])
+    stop_btn.click(stop_generation, None, None)
+
+    demo.launch(share=True)