Update app.py

app.py

@@ -6,9 +6,12 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStream
 import torch
 
 # Carregar modelo local
-model_id = "lambdaindie/lambda-1v-1B"  # Substitua se quiser
+model_id = "lambdaindie/lambda-1v-1B"
 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 = 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()
 
@@ -45,10 +48,7 @@ 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",
@@ -59,21 +59,21 @@ theme = gr.themes.Base(
     block_title_text_color="#fff"
 )
 
-# Flag para parar
+# Flag de parada
 stop_signal = False
 
 def stop_stream():
     global stop_signal
     stop_signal = True
 
-def respond(message, history, system_message, max_tokens, temperature, top_p):
+def respond(history, system_message, max_tokens, temperature, top_p):
     global stop_signal
     stop_signal = False
 
-    # Construção do prompt
+    # Construir prompt
     prompt = ""
     if system_message:
-        prompt += f"{system_message}\n\n"
+        prompt += system_message + "\n\n"
 
     for msg in history:
         role = msg["role"]
@@ -83,11 +83,11 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         elif role == "assistant":
             prompt += f"Assistant: {content}\n"
 
-    prompt += f"User: {message}\nAssistant:"
+    prompt += "Assistant:"
 
     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,
@@ -107,14 +107,18 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         if stop_signal:
             break
         output += token
-        yield {"role": "assistant", "content": output}
+        yield history + [{"role": "assistant", "content": output}]
 
     end = time.time()
-    yield {"role": "system", "content": f"Pensou por {end - start:.1f} segundos"}
+    yield history + [
+        {"role": "assistant", "content": output},
+        {"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")
+    state = gr.State([])
 
     with gr.Row():
         msg = gr.Textbox(label="Mensagem")
@@ -127,16 +131,14 @@ with gr.Blocks(css=css, theme=theme) as app:
         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):
+    def handle_user_msg(user_msg, chat_history):
         if user_msg:
             chat_history = chat_history + [{"role": "user", "content": user_msg}]
         return "", chat_history
 
-    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)
+    send_btn.click(fn=handle_user_msg, inputs=[msg, state], outputs=[msg, state])\
+        .then(fn=respond, inputs=[state, system_message, max_tokens, temperature, top_p], outputs=[chatbot, state])
 
     stop_btn.click(fn=stop_stream, inputs=[], outputs=[])
 
-app.launch(share=True)
+app.launch(share=True