update app

app.py

@@ -73,31 +73,58 @@ def generate_stream(prompt, max_new_tokens=128):
     # thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
     # thread.start()
     
-    inputs = {k: v.to(model.device) for k, v in inputs.items()}
-    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    def run_generation():
-        try:
-            model.generate(
-                **inputs,
-                max_new_tokens=max_new_tokens,
-                do_sample=True,
-                temperature=0.7,
-                top_p=0.9,
-                streamer=streamer,
-                eos_token_id=tokenizer.eos_token_id,
-                pad_token_id=tokenizer.pad_token_id,
-                use_cache=True,
-            )
-        except Exception as e:
-            # simpan error agar bisa dikembalikan ke client setelah streamer selesai
-            # error_container.append(str(e))
-            pass
-
-    thread = threading.Thread(target=run_generation, daemon=True)
-    thread.start()
-
-    for token in streamer:
-        yield token
+    # inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    # streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+    # def run_generation():
+    #     try:
+    #         model.generate(
+    #             **inputs,
+    #             max_new_tokens=max_new_tokens,
+    #             do_sample=True,
+    #             temperature=0.7,
+    #             top_p=0.9,
+    #             streamer=streamer,
+    #             eos_token_id=tokenizer.eos_token_id,
+    #             pad_token_id=tokenizer.pad_token_id,
+    #             use_cache=True,
+    #         )
+    #     except Exception as e:
+    #         # simpan error agar bisa dikembalikan ke client setelah streamer selesai
+    #         # error_container.append(str(e))
+    #         pass
+
+    # thread = threading.Thread(target=run_generation, daemon=True)
+    # thread.start()
+
+    # for token in streamer:
+    #     yield token
+
+
+    streamer = tokenizer.as_target_tokenizer()
+
+    with torch.no_grad():
+        output_ids = model.generate(
+            **inputs,
+            max_new_tokens=128,       # batasi jawaban
+            min_new_tokens=16,        # biar ga berhenti terlalu cepat
+            temperature=0.7,          # lebih to the point
+            top_p=0.9,
+            do_sample=True,
+            early_stopping=True,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+
+    decoded = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    if "Assistant:" in decoded:
+        answer = decoded.split("Assistant:")[-1].strip()
+    else:
+        answer = decoded
+
+    # stream potongan kalimat (kata demi kata)
+    for word in answer.split():
+        yield word + " "
+
 
 @app.post("/stream")
 async def chat(req: ChatRequest):