Update app.py

app.py

@@ -1,44 +1,79 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import torch
+import threading
 
 model_name = "lambdaindie/lambda-1v-1B"
 
+# Carrega modelo na CPU de forma mais leve
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    torch_dtype="float32",
-    low_cpu_mem_usage=True,
-    device_map="auto"
+    torch_dtype=torch.float16,  # ou torch.bfloat16 se suportar
+    low_cpu_mem_usage=True
 )
-model.eval()
-
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-def respond(prompt):
-    full_prompt = f"Think step-by-step.\nQuestion: {prompt}\nAnswer:"
-    inputs = tokenizer(full_prompt, return_tensors="pt", return_attention_mask=False)
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+stop_flag = {"stop": False}
+
+def respond(prompt, history):
+    stop_flag["stop"] = False
+    history = history[-3:]  # Mantém apenas os últimos 3 pares
+
+    full_prompt = f"\nThink a bit step-by-step before answering.\nQuestion: {prompt}\nAnswer:"
+    inputs = tokenizer(full_prompt, return_tensors="pt")
+
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
-    output = model.generate(
-        **inputs,
-        max_new_tokens=128,
-        do_sample=False,  # greedy, menos RAM
-        pad_token_id=tokenizer.eos_token_id,
+    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,
+        }
     )
-    answer = tokenizer.decode(output[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
-    return answer.strip()
+    generation_thread.start()
 
+    reasoning = ""
+    for new_text in streamer:
+        if stop_flag["stop"]:
+            return "", history
+        reasoning += new_text
+        yield "", history + [(prompt, f"<div class='final-answer'>{reasoning}</div>")]
+
+def stop_generation():
+    stop_flag["stop"] = True
+
+# Interface Gradio
 with gr.Blocks(css="""
-    .gr-button, .gr-textbox {  
+    #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;  
+        white-space: pre-wrap;  
+        font-size: 11px !important;  
     }
 """) as demo:
-    gr.Markdown("## λambdAI — Light CPU Reasoning")
-    txt = gr.Textbox(placeholder="Digite sua pergunta...", show_label=False)
-    output = gr.Textbox(label="Resposta", lines=6)
-    btn = gr.Button("Enviar")
+    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")
 
-    btn.click(respond, txt, output)
-    txt.submit(respond, txt, output)
+    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)
+    demo.launch(share=True)