Update app.py

app.py

@@ -6,6 +6,9 @@ import torch
 model_name = "Qwen/Qwen2.5-0.5B"
 try:
     tokenizer = AutoTokenizer.from_pretrained(model_name)
+    # Đặt pad_token_id nếu chưa có
+    if tokenizer.pad_token_id is None:
+        tokenizer.pad_token_id = tokenizer.eos_token_id
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         torch_dtype="auto",
@@ -18,21 +21,29 @@ except Exception as e:
     raise
 
 # Hàm sinh văn bản (dùng cho cả UI và API)
-def generate_text(prompt, max_length=100):
+def generate_text(prompt, max_length, state):
     try:
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        # Mã hóa đầu vào với attention_mask
+        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(model.device)
         outputs = model.generate(
-            inputs["input_ids"],
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
             max_length=max_length,
             num_return_sequences=1,
             no_repeat_ngram_size=2,
             do_sample=True,
             top_k=50,
-            top_p=0.95
+            top_p=0.95,
+            pad_token_id=tokenizer.pad_token_id
         )
-        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Cập nhật state với kết quả mới
+        state.append(generated_text)
+        return state, generated_text  # Trả về state và output để hiển thị
     except Exception as e:
-        return f"Error: {str(e)}"
+        error_msg = f"Error: {str(e)}"
+        state.append(error_msg)
+        return state, error_msg
 
 # Hàm hiển thị thông tin API
 def get_api_info():
@@ -56,6 +67,9 @@ with gr.Blocks(title="Qwen2.5-0.5B Text Generator") as demo:
     gr.Markdown("# Qwen2.5-0.5B Text Generator")
     gr.Markdown("Enter a prompt below or use the API!")
     
+    # State để lưu trữ lịch sử kết quả
+    state = gr.State(value=[])  # Khởi tạo state là danh sách rỗng
+    
     # Hiển thị thông tin API
     gr.Markdown("### API Information")
     api_info = gr.Textbox(label="API Details", value=get_api_info(), interactive=False)
@@ -67,29 +81,29 @@ with gr.Blocks(title="Qwen2.5-0.5B Text Generator") as demo:
         max_length_input = gr.Slider(50, 500, value=100, step=10, label="Max Length")
     
     generate_button = gr.Button("Generate")
-    output_text = gr.Textbox(label="Generated Text", interactive=False)
+    output_text = gr.Textbox(label="Generated Text History", interactive=False, lines=10)
     
     # Liên kết button với hàm generate_text
     generate_button.click(
         fn=generate_text,
-        inputs=[prompt_input, max_length_input],
-        outputs=output_text
+        inputs=[prompt_input, max_length_input, state],
+        outputs=[state, output_text]  # Cập nhật cả state và output_text
     )
 
 # Định nghĩa API endpoints với Gradio
 interface = gr.Interface(
-    fn=generate_text,
+    fn=lambda prompt, max_length: generate_text(prompt, max_length, [])[1],  # Chỉ lấy output, không dùng state cho API
     inputs=["text", "number"],
     outputs="text",
     title="Qwen2.5-0.5B API",
-    api_name="/generate"  # API endpoint: /api/generate
+    api_name="/generate"
 ).queue()
 
 health_interface = gr.Interface(
     fn=health_check,
     inputs=None,
     outputs="text",
-    api_name="/health_check"  # API endpoint: /api/health_check
+    api_name="/health_check"
 )
 
 # Gắn các interface vào demo