adjust generation time

app.py

@@ -2,6 +2,7 @@ import os
 import time
 import warnings
 import re
+import gc
 warnings.filterwarnings("ignore", category=UserWarning, module="torch._utils")
 
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
@@ -16,9 +17,9 @@ def print_system_resources():
     # Get container memory limit (for Docker)
     try:
         with open('/sys/fs/cgroup/memory/memory.limit_in_bytes', 'r') as f:
-            mem_limit = int(f.read().strip()) / 1e9  # Convert to GB
+            mem_limit = min(int(f.read().strip()) / 1e9, 16.0)  # Cap at 16GB for HFS free
     except:
-        mem_limit = 16.0  # Fallback for HFS free (16GB)
+        mem_limit = 16.0  # Fallback for HFS free
     print(f"Total physical memory (psutil): {memory.total/1e9:.2f} GB")
     print(f"Container memory limit: {mem_limit:.2f} GB")
     print(f"CPU usage: {cpu_percent}%")
@@ -67,12 +68,12 @@ def clean_text(text):
     if complete_sentences:
         text = ' '.join(complete_sentences)
     else:
-        # Fallback: Keep until last valid word if no complete sentence
+        # Fallback: Keep until last valid word
         words = text.split()
         text = ' '.join(words[:-1]) if len(words) > 1 else text
     return text
 
-def generate_text(prompt, max_length=50, temperature=0.9):
+def generate_text(prompt, max_length=40, temperature=1.0):
     try:
         start_time = time.time()
         print_system_resources()
@@ -84,18 +85,18 @@ def generate_text(prompt, max_length=50, temperature=0.9):
             truncation=True,
             max_length=max_length
         ).to(device)
+        print(f"Input tokens: {len(inputs['input_ids'][0])}")
         
         # Generate text
         outputs = model.generate(
             input_ids=inputs["input_ids"],
             attention_mask=inputs["attention_mask"],
-            max_new_tokens=15,  # Reduce for speed
-            min_length=10,
+            max_new_tokens=10,  # Reduce for speed
+            min_length=5,
             do_sample=True,
-            top_k=30,  # Reduce for speed
-            top_p=0.9,
+            top_k=20,  # Reduce for speed
+            top_p=0.85,  # Prioritize high-quality tokens
             temperature=temperature,
-            no_repeat_ngram_size=2,
             pad_token_id=tokenizer.pad_token_id
         )
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
@@ -104,6 +105,9 @@ def generate_text(prompt, max_length=50, temperature=0.9):
         cleaned_text = clean_text(generated_text)
         elapsed_time = time.time() - start_time
         print(f"Generation time: {elapsed_time:.2f} seconds")
+        # Clear memory cache
+        gc.collect()
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
         return cleaned_text
     except Exception as e:
         return f"Error generating text: {e}"
@@ -117,8 +121,8 @@ 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(20, 100, value=50, step=10, label="Độ dài tối đa"),
-        gr.Slider(0.7, 1.0, value=0.9, step=0.1, label="Nhiệt độ (Temperature)")
+        gr.Slider(20, 100, value=40, step=10, label="Độ dài tối đa"),
+        gr.Slider(0.7, 1.0, value=1.0, step=0.1, label="Nhiệt độ (Temperature)")
     ],
     outputs="text",
     title="Sinh văn bản tiếng Việt",