Update app.py

app.py

@@ -1,7 +1,8 @@
 import time
+import threading
 import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 
 MODEL_NAME = "daniel-dona/gemma-3-270m-it"
 
@@ -39,16 +40,14 @@ def build_prompt(message, history, system_message, max_ctx_tokens=1024):
                 del msgs[i:i+2]
                 break
 
-def respond(message, history, system_message, max_tokens, temperature, top_p):
+def respond_stream(message, history, system_message, max_tokens, temperature, top_p):
     user_id = "default"  # API bağlarsan burada kullanıcı ID'si ile değiştir
     past = sessions.get(user_id)
 
     if past is None:
-        # İlk mesaj → tüm prompt
         text = build_prompt(message, history, system_message)
         inputs = tokenizer([text], return_tensors="pt").to(model.device)
     else:
-        # Sadece yeni mesajı encode et
         inputs = tokenizer([message], return_tensors="pt").to(model.device)
 
     do_sample = temperature > 0
@@ -61,28 +60,35 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         past_key_values=past
     )
 
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+    thread = threading.Thread(
+        target=model.generate,
+        kwargs={**inputs, **{k: v for k, v in gen_kwargs.items() if v is not None}, "streamer": streamer}
+    )
+
     start_time = time.time()
+    token_count = 0
+
     with torch.inference_mode():
-        outputs = model.generate(**inputs, **{k: v for k, v in gen_kwargs.items() if v is not None},
-                                 return_dict_in_generate=True, output_scores=False)
+        thread.start()
+        for token_text in streamer:
+            token_count += 1
+            yield token_text  # Token anında kullanıcıya akar
+        thread.join()
+
     end_time = time.time()
+    tps = token_count / (end_time - start_time) if (end_time - start_time) > 0 else 0
 
     # KV cache güncelle
-    sessions[user_id] = outputs.past_key_values
-
-    # Yanıtı decode et
-    new_tokens = outputs.sequences[0][inputs["input_ids"].shape[1]:]
-    content = tokenizer.decode(new_tokens, skip_special_tokens=True).strip("\n")
-
-    # T/S hesapla
-    token_count = len(new_tokens)
-    elapsed = end_time - start_time
-    tps = token_count / elapsed if elapsed > 0 else 0
+    # generate() ile streamer kullanıldığında past_key_values doğrudan dönmez,
+    # bu yüzden cache'i burada güncellemek için model.forward tabanlı bir yapı kurmak gerekir.
+    # Basitlik için bu örnekte cache ilk turdan sonra sıfırlanıyor.
+    sessions[user_id] = None  # İstersen burayı ileri seviye cache yönetimi ile değiştirebilirsin.
 
-    return f"{content}\n\n⚡ **Hız:** {tps:.2f} token/sn"
+    yield f"\n\n⚡ **Hız:** {tps:.2f} token/sn"
 
 demo = gr.ChatInterface(
-    respond,
+    respond_stream,
     additional_inputs=[
         gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=256, step=1, label="Max new tokens"),