Update app.py

app.py

@@ -8,7 +8,7 @@ from threading import Thread
 
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
-MODEL_ID = "kodetr/stunting-qa-v4"
+MODEL_ID = "kodetr/stunting-qa-v5"
 MODELS = os.environ.get("MODELS")
 
 TITLE = "<h1><center>KONSULTASI STUNTING</center></h1>"
@@ -51,32 +51,60 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
 
     print(f"Conversation is -\n{conversation}")
     
-    input_ids = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer(input_ids, return_tensors="pt").to(0) #gpu 0, cpu 1
+    # ------- use model stunting V5 -------
+    terminators = [
+        text_pipeline.tokenizer.eos_token_id,
+        text_pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
+    ]
     
-    streamer = TextIteratorStreamer(tokenizer, timeout=60., skip_prompt=True, skip_special_tokens=True)
-
-    generate_kwargs = dict(
-        inputs, 
-        streamer=streamer,
-        top_k=top_k,
-        top_p=top_p,
-        repetition_penalty=penalty,
-        max_new_tokens=max_new_tokens, 
-        do_sample=True, 
+    # Hasil dari pipeline akan berupa list dengan dictionary berisi text
+    outputs = text_pipeline(
+        conversation,
+        max_new_tokens=max_new_tokens,
+        eos_token_id=terminators,
+        do_sample=True,
         temperature=temperature,
-        pad_token_id=128000,
-        eos_token_id=[128001,128008,128009],
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=penalty
     )
-    
-    thread = Thread(target=model.generate, kwargs=generate_kwargs)
-    thread.start()
 
+    # Karena pipeline tidak support streaming per token, kita bisa stream per kalimat atau per paragraf
+    full_text = outputs[0]["generated_text"]
     buffer = ""
-    for new_text in streamer:
-        buffer += new_text
+    for part in full_text.split(". "):  # Stream berdasarkan kalimat
+        buffer += part.strip() + ". "
         yield buffer
 
+    # ------- use model stunting V6 -------
+    
+    
+    # input_ids = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)
+    # inputs = tokenizer(input_ids, return_tensors="pt").to(0) #gpu 0, cpu 1
+    
+    # streamer = TextIteratorStreamer(tokenizer, timeout=60., skip_prompt=True, skip_special_tokens=True)
+
+    # generate_kwargs = dict(
+    #     inputs, 
+    #     streamer=streamer,
+    #     top_k=top_k,
+    #     top_p=top_p,
+    #     repetition_penalty=penalty,
+    #     max_new_tokens=max_new_tokens, 
+    #     do_sample=True, 
+    #     temperature=temperature,
+    #     pad_token_id=128000,
+    #     eos_token_id=[128001,128008,128009],
+    # )
+    
+    # thread = Thread(target=model.generate, kwargs=generate_kwargs)
+    # thread.start()
+
+    # buffer = ""
+    # for new_text in streamer:
+    #     buffer += new_text
+    #     yield buffer
+
 
 
 chatbot = gr.Chatbot(height=600)