Update app.py

app.py

@@ -1,21 +1,20 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+import torch
+import threading
 import time
 
-client = InferenceClient("lambdaindie/lambdai")
+model_id = "lambdaindie/lambdai"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
 css = """
 @import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono&display=swap');
-
-* {
-    font-family: 'JetBrains Mono', monospace !important;
-}
-
-body {
-    background-color: #111;
-    color: #e0e0e0;
-}
-
+* { font-family: 'JetBrains Mono', monospace !important; }
+body { background-color: #111; color: #e0e0e0; }
 .markdown-think {
     background-color: #1e1e1e;
     border-left: 4px solid #555;
@@ -25,7 +24,6 @@ body {
     white-space: pre-wrap;
     animation: pulse 1.5s infinite ease-in-out;
 }
-
 @keyframes pulse {
     0% { opacity: 0.6; }
     50% { opacity: 1.0; }
@@ -42,54 +40,55 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         if assistant:
             messages.append({"role": "assistant", "content": assistant})
 
-    thinking_prompt = messages + [{
-        "role": "user",
-        "content": f"{message}\n\nThink a bit step-by-step before answering."
-    }]
+    thinking_prompt = messages + [{"role": "user", "content": f"{message}\n\nThink step-by-step."}]
+    prompt = tokenizer.apply_chat_template(thinking_prompt, tokenize=False, add_generation_prompt=True)
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
 
     reasoning = ""
     yield '<div class="markdown-think">Thinking...</div>'
-
     start = time.time()
 
-    for chunk in client.chat_completion(
-        thinking_prompt,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = chunk.choices[0].delta.content or ""
+    thread = threading.Thread(target=model.generate, kwargs={
+        "inputs": inputs["input_ids"],
+        "max_new_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "streamer": streamer,
+    })
+    thread.start()
+
+    for token in streamer:
         reasoning += token
-        styled_thought = f'<div class="markdown-think">{reasoning.strip()}</div>'
-        yield styled_thought
+        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>"""
 
-    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>
-    """
-
-    time.sleep(2)
-
+    # Segunda etapa: resposta final
     final_prompt = messages + [
         {"role": "user", "content": message},
         {"role": "assistant", "content": reasoning.strip()},
-        {"role": "user", "content": "Now answer based on your reasoning above."}
+        {"role": "user", "content": "Agora responda baseado nisso."}
     ]
+    prompt2 = tokenizer.apply_chat_template(final_prompt, tokenize=False, add_generation_prompt=True)
+    inputs2 = tokenizer(prompt2, return_tensors="pt").to(device)
+    streamer2 = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    thread2 = threading.Thread(target=model.generate, kwargs={
+        "inputs": inputs2["input_ids"],
+        "max_new_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "streamer": streamer2,
+    })
+    thread2.start()
 
     final_answer = ""
-    for chunk in client.chat_completion(
-        final_prompt,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = chunk.choices[0].delta.content or ""
+    for token in streamer2:
         final_answer += token
         yield final_answer.strip()
 
@@ -99,8 +98,7 @@ demo = gr.ChatInterface(
     theme=gr.themes.Base(),
     css=css,
     additional_inputs=[
-        gr.Textbox(value="",
-                   label="System Message"),
+        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")