adjust generation time

app.py

@@ -3,6 +3,7 @@ import time
 import warnings
 import re
 import gc
+import tracemalloc
 warnings.filterwarnings("ignore", category=UserWarning, module="torch._utils")
 
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
@@ -59,16 +60,18 @@ def clean_text(text):
     text = re.sub(r'\s+', ' ', text).strip()
     return text
 
-def generate_text(prompt, temperature=0.9, max_new_tokens=100, top_k=30, max_length=100):
+def generate_text(prompt, temperature=0.7, max_new_tokens=50, min_length=50, max_length=100):
     try:
         start_time = time.time()
+        # Start tracemalloc for memory debugging
+        tracemalloc.start()
         # Debug memory before generation
         print("Memory before generation:")
         print_system_resources()
         # Fixed parameters
-        repetition_penalty = 1.0
+        top_k = 20
+        repetition_penalty = 1.5
         top_p = 0.9
-        min_length = 50  # Force longer output
         # Log parameters
         print(f"Parameters: max_length={max_length}, temperature={temperature}, max_new_tokens={max_new_tokens}, top_k={top_k}, top_p={top_p}, repetition_penalty={repetition_penalty}, min_length={min_length}")
         # Encode input with attention mask
@@ -82,8 +85,8 @@ def generate_text(prompt, temperature=0.9, max_new_tokens=100, top_k=30, max_len
         ).to(device)
         print(f"Encoding time: {time.time() - encode_time:.2f} seconds")
         print(f"Input tokens: {len(inputs['input_ids'][0])}")
-        # Avoid stopping at '.', '!', '?' to allow longer text
-        eos_token_ids = []  # Disable EOS tokens
+        # 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()
@@ -91,14 +94,14 @@ def generate_text(prompt, temperature=0.9, max_new_tokens=100, top_k=30, max_len
             input_ids=inputs["input_ids"],
             attention_mask=inputs["attention_mask"],
             max_new_tokens=int(max_new_tokens),
-            min_length=min_length,
+            min_length=int(min_length),
             do_sample=True,
-            top_k=int(top_k),
+            top_k=top_k,
             top_p=top_p,
             temperature=temperature,
             repetition_penalty=repetition_penalty,
             pad_token_id=tokenizer.pad_token_id,
-            eos_token_id=eos_token_ids if eos_token_ids else None
+            eos_token_id=eos_token_ids
         )
         gen_duration = time.time() - gen_time
         print(f"Generation time: {gen_duration:.2f} seconds")
@@ -110,6 +113,12 @@ def generate_text(prompt, temperature=0.9, max_new_tokens=100, top_k=30, max_len
         print(f"Cleaned output: {cleaned_text}")
         elapsed_time = time.time() - start_time
         print(f"Total time: {elapsed_time:.2f} seconds")
+        # Check for repetitive text
+        from collections import Counter
+        words = cleaned_text.split()
+        word_counts = Counter(words)
+        if any(count > 3 for count in word_counts.values()):
+            print("Warning: Repetitive text detected")
         # Clear memory cache
         gc.collect()
         if torch.cuda.is_available():
@@ -117,6 +126,11 @@ def generate_text(prompt, temperature=0.9, max_new_tokens=100, top_k=30, max_len
         # Debug memory after generation
         print("Memory after generation:")
         print_system_resources()
+        # Log tracemalloc
+        snapshot = tracemalloc.take_snapshot()
+        top_stats = snapshot.statistics('lineno')
+        print("[Top 5 memory]", top_stats[:5])
+        tracemalloc.stop()
         return cleaned_text
     except Exception as e:
         return f"Error generating text: {e}"
@@ -128,16 +142,16 @@ demo = gr.Interface(
         gr.Textbox(
             label="Nhập văn bản đầu vào",
             placeholder="Viết gì đó bằng tiếng Việt...",
-            value="Hôm nay là một ngày đẹp trời, hãy kể một câu chuyện dài về trải nghiệm của bạn."  # Updated prompt
+            value="Hôm nay là một ngày đẹp trời, hãy kể một câu chuyện dài về một trải nghiệm đáng nhớ của bạn, bao gồm chi tiết về bối cảnh, nhân vật, và cảm xúc."
         ),
-        gr.Slider(0.1, 0.9, value=0.9, step=0.1, label="Nhiệt độ (Temperature, 0.1-0.3 cho mạch lạc, 0.7-0.9 cho câu dài và sáng tạo)"),
-        gr.Slider(25, 100, value=100, step=5, label="Số token mới tối đa (max_new_tokens, 25-50 cho câu ngắn, 75-100 cho câu dài)"),
-        gr.Slider(5, 30, value=30, step=5, label="Top K (top_k, 5-10 cho mạch lạc, 20-30 cho đa dạng và dài hơn)"),
-        gr.Slider(30, 100, value=100, step=5, label="Độ dài tối đa (max_length, 30-50 cho câu ngắn, 75-100 cho câu dài)")
+        gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Nhiệt độ (Temperature, 0.1-0.3 cho mạch lạc, 0.7-1.5 cho sáng tạo và câu dài)"),
+        gr.Slider(25, 100, value=50, step=5, label="Số token mới tối đa (max_new_tokens, 25-50 cho câu ngắn, 75-100 cho câu dài)"),
+        gr.Slider(30, 100, value=50, step=5, label="Độ dài tối thiểu (min_length, 30-50 cho câu vừa, 75-100 cho câu dài)"),
+        gr.Slider(50, 150, value=100, step=5, label="Độ dài tối đa (max_length, 50-100 cho câu ngắn, 100-150 cho câu dài)")
     ],
     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. Chọn temperature 0.7-0.9, max_new_tokens 75-100, top_k 20-30, và max_length 75-100 để tạo câu dài và đa dạng. Dùng temperature 0.1-0.3, max_new_tokens 25-50, top_k 5-10, và max_length 30-50 để đạt thời gian <2 giây và văn bản mạch lạc. Lưu ý: Để tạo câu dài, dùng prompt yêu cầu câu chuyện hoặc mô tả chi tiết, và thử temperature=0.9, top_k=30, min_length=50.",
+    description="Dùng mô hình GPT-2 Vietnamese từ NlpHUST để sinh văn bản tiếng Việt. Chọn temperature 0.7, max_new_tokens 50, min_length 50, và max_length 100 để tạo câu dài, mạch lạc, và hoàn chỉnh. Dùng temperature 0.1-0.3, max_new_tokens 25-50, min_length 30-50 cho câu ngắn và nhanh. Lưu ý: Dùng prompt chi tiết (bối cảnh, nhân vật, cảm xúc) để cải thiện chất lượng câu chuyện dài. Tránh max_new_tokens=100 với temperature=0.3 để giảm lặp từ. Top_k mặc định là 20.",
     allow_flagging="never"
 )