improve response time

app.py

@@ -1,6 +1,7 @@
 import os
 import time
 import warnings
+import re
 warnings.filterwarnings("ignore", category=UserWarning, module="torch._utils")
 
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
@@ -8,6 +9,13 @@ import torch
 import gradio as gr
 import psutil
 
+# Print system resources for debugging
+def print_system_resources():
+    cpu_percent = psutil.cpu_percent(interval=1)
+    memory = psutil.virtual_memory()
+    print(f"CPU usage: {cpu_percent}%")
+    print(f"Memory usage: {memory.percent}% ({memory.used/1e9:.2f}/{memory.total/1e9:.2f} GB)")
+
 # Load model and tokenizer
 model_id = "NlpHUST/gpt2-vietnamese"
 try:
@@ -28,23 +36,20 @@ model.to(device)
 model.eval()
 
 # Print device and memory info for debugging
-print(f"---------- Info -----------")
 print(f"Device: {device}")
 print(f"Memory allocated: {torch.cuda.memory_allocated(device)/1e9:.2f} GB" if torch.cuda.is_available() else "CPU only")
-
-def print_system_resources():
-    cpu_percent = psutil.cpu_percent(interval=1)
-    memory = psutil.virtual_memory()
-    print(f"CPU usage: {cpu_percent}%")
-    print(f"Memory usage: {memory.percent}% ({memory.used/1e9:.2f}/{memory.total/1e9:.2f} GB)")
-
-# Call before generation
 print_system_resources()
-print(f"--------------------------")
 
-def generate_text(prompt, max_length=50, temperature=1.0):
+def clean_text(text):
+    """Clean generated text by removing non-alphabetic characters and extra spaces."""
+    text = re.sub(r'[^\w\s.,!?]', '', text)  # Remove non-alphabetic characters
+    text = re.sub(r'\s+', ' ', text).strip()  # Normalize spaces
+    return text
+
+def generate_text(prompt, max_length=50, temperature=0.7):
     try:
         start_time = time.time()
+        print_system_resources()  # Print resources before generation
         # Encode input with attention mask
         inputs = tokenizer(
             prompt,
@@ -58,18 +63,21 @@ def generate_text(prompt, max_length=50, temperature=1.0):
         outputs = model.generate(
             input_ids=inputs["input_ids"],
             attention_mask=inputs["attention_mask"],
-            max_new_tokens=30,  # Limit new tokens to reduce computation
+            max_new_tokens=30,
+            min_length=10,  # Ensure minimum output length
             temperature=temperature,
-            do_sample=True,
-            num_beams=3,  # Use beam search for faster generation
-            no_repeat_ngram_size=2,  # Prevent repetitive phrases
+            do_sample=False,  # Use greedy decoding for consistency
+            num_beams=3,  # Use beam search for better quality
+            no_repeat_ngram_size=2,
             pad_token_id=tokenizer.pad_token_id,
-            early_stopping=True  # Stop when generation is complete
+            early_stopping=True
         )
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        cleaned_text = clean_text(generated_text)
         elapsed_time = time.time() - start_time
+        print_system_resources()
         print(f"Generation time: {elapsed_time:.2f} seconds")
-        return generated_text
+        return cleaned_text
     except Exception as e:
         return f"Error generating text: {e}"
 
@@ -79,7 +87,7 @@ demo = gr.Interface(
     inputs=[
         gr.Textbox(label="Nhập văn bản đầu vào", placeholder="Viết gì đó bằng tiếng Việt..."),
         gr.Slider(20, 100, value=50, step=10, label="Độ dài tối đa"),
-        gr.Slider(0.5, 1.5, value=1.0, step=0.1, label="Nhiệt độ (Temperature)")
+        gr.Slider(0.5, 1.0, value=0.7, step=0.1, label="Nhiệt độ (Temperature)")
     ],
     outputs="text",
     title="Sinh văn bản tiếng Việt",