Update app.py

app.py

@@ -1,76 +1,44 @@
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-import threading
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Detectar dispositivo automaticamente
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# Inicializar o modelo e o 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}
 
-def respond(prompt, history):
-    stop_flag["stop"] = False
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype="float32",
+    low_cpu_mem_usage=True,
+    device_map="auto"
+)
+model.eval()
 
-    full_prompt = f"\nThink a bit step-by-step before answering.  \nQuestion: {prompt} \nAnswer:"
-    inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+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()}
 
-    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,
-        }
+    output = model.generate(
+        **inputs,
+        max_new_tokens=128,
+        do_sample=False,  # greedy, menos RAM
+        pad_token_id=tokenizer.eos_token_id,
     )
-    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
+    answer = tokenizer.decode(output[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
+    return answer.strip()
 
 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;  
+    .gr-button, .gr-textbox {  
         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")
+    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")
 
-    send_btn.click(respond, [txt, chatbot], [txt, chatbot])
-    txt.submit(respond, [txt, chatbot], [txt, chatbot])
-    stop_btn.click(stop_generation, None, None)
+    btn.click(respond, txt, output)
+    txt.submit(respond, txt, output)
 
-    demo.launch(share=True)
+demo.launch(share=True)