Update app.py

app.py

@@ -1,88 +1,23 @@
-import gradio as gr
-import torch
+import os
 import time
 import threading
+import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+import torch
 
-# === Carregar modelo local (CPU) ===
-model_name = "lambdaindie/lambda-1v-1B"  # troque pelo teu
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name).to("cpu")  # <- CPU aqui
-
-# === Streamer global para interrupção ===
-stop_signal = {"stop": False}
-
-def generate_stream(prompt, max_tokens=512, temperature=0.7, top_p=0.95):
-    stop_signal["stop"] = False
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
-
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-
-    generation_thread = threading.Thread(
-        target=model.generate,
-        kwargs=dict(
-            input_ids=inputs["input_ids"],
-            attention_mask=inputs["attention_mask"],
-            streamer=streamer,
-            max_new_tokens=max_tokens,
-            do_sample=True,
-            temperature=temperature,
-            top_p=top_p,
-            pad_token_id=tokenizer.eos_token_id,
-        )
-    )
-    generation_thread.start()
-
-    output = ""
-    for token in streamer:
-        if stop_signal["stop"]:
-            break
-        output += token
-        yield output.strip()
-
-def stop_stream():
-    stop_signal["stop"] = True
-
-def respond(message, history, system_message, max_tokens, temperature, top_p):
-    messages = [{"role": "system", "content": system_message}] if system_message else []
-
-    for user, assistant in history[-3:]:  # Limita a 3 interações passadas
-        if user:
-            messages.append({"role": "user", "content": user})
-        if assistant:
-            messages.append({"role": "assistant", "content": assistant})
-
-    thinking_prompt = messages + [{"role": "user", "content": f"{message}\n\nThink step-by-step before answering."}]
-    thinking_text = "\n".join([f"{m['role']}: {m['content']}" for m in thinking_prompt])
-
-    reasoning = ""
-    yield '<div class="markdown-think">Thinking...</div>'
-
-    start = time.time()
-    for token in generate_stream(thinking_text, max_tokens, temperature, top_p):
-        reasoning = token
-        yield f'<div class="markdown-think">{reasoning.strip()}</div>'
-
-    elapsed = time.time() - start
-    yield f"""
-    <div style="margin-top:12px;padding:8px 12px;background-color:#222;border-left:4px solid #888;
-            font-family:'JetBrains Mono', monospace;color:#ccc;font-size:14px;">
-        Pensou por {elapsed:.1f} segundos
-    </div>
-    """
-
-    final_prompt = thinking_text + f"\n\nuser: {message}\nassistant: {reasoning.strip()}\nuser: Now answer based on your reasoning above.\nassistant:"
-    final_answer = ""
-
-    for token in generate_stream(final_prompt, max_tokens, temperature, top_p):
-        final_answer = token
-        yield final_answer.strip()
-
-# === Interface ===
+# Carregar modelo local
+model_id = "lambdaindie/lambda-1v-1B"  # Substitua se quiser
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+model.to("cuda" if torch.cuda.is_available() else "cpu")
+model.eval()
 
+# Estilo
 css = """
 @import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono&display=swap');
-* { font-family: 'JetBrains Mono', monospace !important; }
+* {
+    font-family: 'JetBrains Mono', monospace !important;
+}
 html, body, .gradio-container {
     background-color: #111 !important;
     color: #e0e0e0 !important;
@@ -110,36 +45,98 @@ textarea, input, button, select {
 
 theme = gr.themes.Base(
     primary_hue="gray",
-    font=[gr.themes.GoogleFont("JetBrains Mono"), "monospace"]
+    font=[
+        gr.themes.GoogleFont("JetBrains Mono"),
+        "monospace"
+    ]
 ).set(
     body_background_fill="#111",
     body_text_color="#e0e0e0",
-    input_background_fill="#222",
-    input_border_color="#444",
     button_primary_background_fill="#333",
     button_primary_text_color="#e0e0e0",
+    input_background_fill="#222",
+    input_border_color="#444",
+    block_title_text_color="#fff"
 )
 
-chatbot = gr.ChatInterface(
-    fn=respond,
-    title="λambdAI",
-    css=css,
-    theme=theme,
-    additional_inputs=[
-        gr.Textbox(value="", label="System Message"),
-        gr.Slider(64, 2048, value=512, step=1, label="Max Tokens"),
-        gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
-    ]
-)
+# Flag para parar
+stop_signal = False
+
+def stop_stream():
+    global stop_signal
+    stop_signal = True
+
+def respond(message, history, system_message, max_tokens, temperature, top_p):
+    global stop_signal
+    stop_signal = False
+
+    # Construção do prompt
+    prompt = ""
+    if system_message:
+        prompt += f"{system_message}\n\n"
+
+    for msg in history:
+        role = msg["role"]
+        content = msg["content"]
+        if role == "user":
+            prompt += f"User: {content}\n"
+        elif role == "assistant":
+            prompt += f"Assistant: {content}\n"
+
+    prompt += f"User: {message}\nAssistant:"
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+    )
+
+    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    output = ""
+    start = time.time()
+
+    for token in streamer:
+        if stop_signal:
+            break
+        output += token
+        yield {"role": "assistant", "content": output}
+
+    end = time.time()
+    yield {"role": "system", "content": f"Pensou por {end - start:.1f} segundos"}
+
+# Interface
+with gr.Blocks(css=css, theme=theme) as app:
+    chatbot = gr.Chatbot(label="λ", type="messages")
+
+    with gr.Row():
+        msg = gr.Textbox(label="Mensagem")
+        send_btn = gr.Button("Enviar")
+        stop_btn = gr.Button("Parar")
+
+    with gr.Accordion("Configurações Avançadas", open=False):
+        system_message = gr.Textbox(label="System Message", value="")
+        max_tokens = gr.Slider(64, 2048, value=256, step=1, label="Max Tokens")
+        temperature = gr.Slider(0.1, 2.0, value=0.7, step=0.1, label="Temperature")
+        top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
+
+    state = gr.State([])
+
+    def user_message_submit(user_msg, chat_history):
+        if user_msg:
+            chat_history = chat_history + [{"role": "user", "content": user_msg}]
+        return "", chat_history
 
-stop_btn = gr.Button("Parar Geração")
-stop_btn.click(fn=stop_stream, inputs=[], outputs=[])
+    send_btn.click(fn=user_message_submit, inputs=[msg, state], outputs=[msg, state])\
+        .then(fn=respond, inputs=[msg, state, system_message, max_tokens, temperature, top_p], outputs=chatbot)
 
-app = gr.Blocks()
-with app:
-    chatbot.render()
-    stop_btn.render()
+    stop_btn.click(fn=stop_stream, inputs=[], outputs=[])
 
-if __name__ == "__main__":
-    app.launch(share=True)
+app.launch(share=True)