adjust generation time

app.py

@@ -3,7 +3,7 @@ import time
 import warnings
 import re
 import gc
-warnings.filterwarnings("ignore", category=UserWarning, module="torch._utils")
+warnings.filterWarnings("ignore", category=UserWarning, module="torch._utils")
 
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import torch
@@ -65,15 +65,18 @@ def clean_text(text):
     text = re.sub(r'\s+', ' ', text).strip()
     return text
 
-def generate_text(prompt, temperature=0.7, max_new_tokens=10):
+def generate_text(prompt, temperature=0.5, max_new_tokens=30):
     try:
         start_time = time.time()
         print_system_resources()
+        # Fixed parameters
+        max_length = 50
+        top_k = 20
+        repetition_penalty = 1.2
         # Log parameters
-        max_length = 50  # Fixed
-        top_k = 20       # Fixed
-        print(f"Parameters: max_length={max_length}, temperature={temperature}, max_new_tokens={max_new_tokens}, top_k={top_k}")
+        print(f"Parameters: max_length={max_length}, temperature={temperature}, max_new_tokens={max_new_tokens}, top_k={top_k}, repetition_penalty={repetition_penalty}")
         # Encode input with attention mask
+        encode_time = time.time()
         inputs = tokenizer(
             prompt,
             return_tensors="pt",
@@ -81,9 +84,13 @@ def generate_text(prompt, temperature=0.7, max_new_tokens=10):
             truncation=True,
             max_length=max_length
         ).to(device)
+        print(f"Encoding time: {time.time() - encode_time:.2f} seconds")
         print(f"Input tokens: {len(inputs['input_ids'][0])}")
-        
+        # Define EOS token IDs for '.', '!', '?'
+        eos_token_ids = [tokenizer.encode(s)[0] for s in ['.', '!', '?']]
+        print(f"EOS token IDs: {eos_token_ids}")
         # Generate text
+        gen_time = time.time()
         outputs = model.generate(
             input_ids=inputs["input_ids"],
             attention_mask=inputs["attention_mask"],
@@ -92,16 +99,18 @@ def generate_text(prompt, temperature=0.7, max_new_tokens=10):
             do_sample=True,
             top_k=int(top_k),
             temperature=temperature,
+            repetition_penalty=repetition_penalty,
             pad_token_id=tokenizer.pad_token_id,
-            eos_token_id=tokenizer.encode('.')[0]  # Stop at period
+            eos_token_id=eos_token_ids
         )
+        print(f"Generation time: {time.time() - gen_time:.2f} seconds")
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         print(f"Raw output: {generated_text}")
         print(f"Generated token count: {len(outputs[0])}")
         cleaned_text = clean_text(generated_text)
         print(f"Cleaned output: {cleaned_text}")
         elapsed_time = time.time() - start_time
-        print(f"Generation time: {elapsed_time:.2f} seconds")
+        print(f"Total time: {elapsed_time:.2f} seconds")
         # Clear memory cache
         gc.collect()
         return cleaned_text
@@ -117,12 +126,12 @@ demo = gr.Interface(
             placeholder="Viết gì đó bằng tiếng Việt...",
             value="Hôm nay là một ngày đẹp trời"  # Default text
         ),
-        gr.Slider(0.5, 1.5, value=0.7, step=0.1, label="Nhiệt độ (Temperature)"),
-        gr.Slider(5, 30, value=10, step=1, label="Số token mới tối đa (max_new_tokens)")
+        gr.Slider(0.3, 0.7, value=0.5, step=0.1, label="Nhiệt độ (Temperature, 0.3-0.5 cho tốc độ nhanh, 0.6-0.7 cho đa dạng hơn)"),
+        gr.Slider(20, 50, value=30, step=5, label="Số token mới tối đa (max_new_tokens, 20-30 cho tốc độ nhanh, 40-50 cho câu dài hơn)")
     ],
     outputs="text",
     title="Sinh văn bản tiếng Việt",
-    description="Dùng mô hình GPT-2 Vietnamese từ NlpHUST để sinh văn bản tiếng Việt. Điều chỉnh temperature và max_new_tokens để tối ưu thời gian và chất lượng đầu ra.",
+    description="Dùng mô hình GPT-2 Vietnamese từ NlpHUST để sinh văn bản tiếng Việt. Chọn temperature 0.3-0.5 và max_new_tokens 20-30 để đạt thời gian <2 giây. Dùng temperature 0.6-0.7 và max_new_tokens 40-50 cho câu dài và đa dạng hơn.",
     allow_flagging="never"
 )