Update app.py

app.py

@@ -38,22 +38,23 @@ h3 {
 # ------- use model stunting V5 -------
 # -------------------------------------
 
-text_pipeline = pipeline(
-    "text-generation",
-    model=MODEL_ID,
-    model_kwargs={"torch_dtype": torch.bfloat16},
-    device_map="auto",
-)
+# text_pipeline = pipeline(
+#     "text-generation",
+#     model=MODEL_ID,
+#     model_kwargs={"torch_dtype": torch.bfloat16},
+#     device_map="auto",
+# )
+
 # -------------------------------------
 # ------- use model stunting V6 -------
 # -------------------------------------
 
-# model = AutoModelForCausalLM.from_pretrained(
-#           MODEL_ID,
-#           torch_dtype=torch.bfloat16,
-#           device_map="auto",
-#         )
-# tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+model = AutoModelForCausalLM.from_pretrained(
+          MODEL_ID,
+          torch_dtype=torch.bfloat16,
+          device_map="auto",
+        )
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
 @spaces.GPU
 def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int, top_p: float, top_k: int, penalty: float):
@@ -71,72 +72,72 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
 # -------------------------------------
 
     # Ubah ke format prompt-style string
-    conversation_text = ""
-    for turn in conversation:
-        role = turn["role"]
-        content = turn["content"]
-        if role == "system":
-            conversation_text += f"[SYSTEM]: {content}\n"
-        elif role == "user":
-            conversation_text += f"[USER]: {content}\n"
-        elif role == "assistant":
-            conversation_text += f"[ASSISTANT]: {content}\n"
+    # conversation_text = ""
+    # for turn in conversation:
+    #     role = turn["role"]
+    #     content = turn["content"]
+    #     if role == "system":
+    #         conversation_text += f"[SYSTEM]: {content}\n"
+    #     elif role == "user":
+    #         conversation_text += f"[USER]: {content}\n"
+    #     elif role == "assistant":
+    #         conversation_text += f"[ASSISTANT]: {content}\n"
     
-    terminators = [
-        text_pipeline.tokenizer.eos_token_id,
-        text_pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
-    ]
+    # terminators = [
+    #     text_pipeline.tokenizer.eos_token_id,
+    #     text_pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
+    # ]
     
     # Hasil dari pipeline akan berupa list dengan dictionary berisi text
-    outputs = text_pipeline(
-        conversation_text,
-        max_new_tokens=max_new_tokens,
-        eos_token_id=terminators,
-        do_sample=True,
-        temperature=temperature,
-        top_p=top_p,
-        top_k=top_k,
-        repetition_penalty=penalty
-    )
+    # outputs = text_pipeline(
+    #     conversation_text,
+    #     max_new_tokens=max_new_tokens,
+    #     eos_token_id=terminators,
+    #     do_sample=True,
+    #     temperature=temperature,
+    #     top_p=top_p,
+    #     top_k=top_k,
+    #     repetition_penalty=penalty
+    # )
 
     # 4. Ekstrak teks hasil dan stream per kalimat
-    generated_text = outputs[0].get("generated_text", "")
-    streamed_text = generated_text[len(conversation_text):].strip()  # Hilangkan prompt awal
+    # generated_text = outputs[0].get("generated_text", "")
+    # streamed_text = generated_text[len(conversation_text):].strip()  # Hilangkan prompt awal
 
-    buffer = ""
-    for part in streamed_text.split(". "):
-        buffer += part.strip() + ". "
-        yield buffer
+    # buffer = ""
+    # for part in streamed_text.split(". "):
+    #     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
+    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)
+    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],
-    # )
+    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()
+    thread = Thread(target=model.generate, kwargs=generate_kwargs)
+    thread.start()
 
-    # buffer = ""
-    # for new_text in streamer:
-    #     buffer += new_text
-    #     yield buffer
+    buffer = ""
+    for new_text in streamer:
+        buffer += new_text
+        yield buffer