modify app.py for load models pipeline function

app.py

@@ -2,7 +2,7 @@
 import gradio as gr
 import spacy
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline, AutoConfig
-from transformers import ProphetNetTokenizer, ProphetNetForConditionalGeneration, T5Tokenizer, T5ForConditionalGeneration
+from transformers import ProphetNetTokenizer, ProphetNetForConditionalGeneration, T5Tokenizer, T5ForConditionalGeneration, BartTokenizer, BartForConditionalGeneration
 import torch
 import time
 import re
@@ -16,8 +16,7 @@ nlp = spacy.load("en_core_web_md")
 
 # Đường dẫn mô hình
 MODEL_PATHS = {
-    "prophetnet-finetuned": "ManB2207540/prophetnet_SQuAD_1.1-2epoch_break",
-    "prophetnet tieu chuan": "microsoft/prophetnet-large-uncased-squad-qg",
+    "prophetnet-large-uncased-finetuned": "ManB2207540/prophetnet_SQuAD_1.1-2epoch_break",
     "bart-finetuned": "mghan3624/bart_qg_finetune_squad",
     "t5-small-finetuned": "tbtminh/t5-small-qg-finetuned"
 }
@@ -39,14 +38,48 @@ def load_t5_pipeline(model_path):
         print(f"Failed to load T5 pipeline for {model_path}: {e}")
         return None
 
-# Hàm tải mô hình chung
-def load_pipeline(model_path):
+# Ham tải mô hình ProphetNet
+def load_prophetnet_pipeline(model_path):
+    try:
+        tokenizer = ProphetNetTokenizer.from_pretrained(model_path)
+        model = ProphetNetForConditionalGeneration.from_pretrained(model_path)
+        return pipeline(
+            "text2text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_length=256,
+            num_return_sequences=1,
+            device=0 if torch.cuda.is_available() else -1
+        )
+    except Exception as e:
+        print(f"Failed to load ProphetNet pipeline for {model_path}: {e}")
+        return None
+
+# Hàm tải mô hình Bart
+def load_bart_pipeline(model_path):
     try:
         config = AutoConfig.from_pretrained(model_path)
         if getattr(config, "early_stopping", None) is None:
             config.early_stopping = False
+        tokenizer = BartTokenizer.from_pretrained(model_path)
+        model = BartForConditionalGeneration.from_pretrained(model_path, config=config)
+        return pipeline(
+            "text2text-generation",
+            model=model,
+            tokenizer=tokenizer,
+            max_length=256,
+            num_return_sequences=1,
+            device=0 if torch.cuda.is_available() else -1
+        )
+    except Exception as e:
+        print(f"Failed to load Bart pipeline for {model_path}: {e}")
+        return None
+
+# Hàm tải mô hình chung
+def load_pipeline(model_path):
+    try:
         tokenizer = AutoTokenizer.from_pretrained(model_path)
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_path, config=config)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
         return pipeline(
             "text2text-generation",
             model=model,
@@ -67,6 +100,10 @@ def get_pipeline(model_name):
     if model_name not in pipeline_cache:
         if model_name == "t5-small-finetuned":
             pipeline_cache[model_name] = load_t5_pipeline(model_path)
+        elif model_name == "prophetnet-large-uncased-finetuned":
+            pipeline_cache[model_name] = load_prophetnet_pipeline(model_path)
+        elif model_name == "bart-finetuned":
+            pipeline_cache[model_name] = load_bart_pipeline(model_path)
         else:
             pipeline_cache[model_name] = load_pipeline(model_path)
     return pipeline_cache[model_name]
@@ -89,6 +126,10 @@ def generate_question(context, answer, model_name):
     tokenizer = pipe.tokenizer
     if model_name == "t5-small-finetuned":
         prompt = f"generate question: context: {context} answer: {answer}"
+    elif model_name == "prophetnet-large-uncased-finetuned":
+        prompt = f"context: {context} answer: {answer}"
+    elif model_name == "bart-finetuned":
+        prompt = f"context: {context} answer: {answer}"
     else:
         prompt = f"context: {context} answer: {answer}"
     print(f"Prompt: {prompt}")  # In ra prompt để kiểm tra
@@ -97,7 +138,7 @@ def generate_question(context, answer, model_name):
     encoded = tokenizer(prompt, return_tensors="pt", truncation=False, max_length=512)
     input_ids = encoded["input_ids"]
     if input_ids.size(1) > 512:
-        return "❌ Context quá dài (hơn 512 token). Xin nhập vào context ngắn hơn."
+        return "❌ Văn bản quá dài. Xin nhập vào văn bản ngắn hơn." # (hơn 512 token)
 
     # Proceed with tokenization (with truncation if needed)
     encoded = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)
@@ -173,11 +214,11 @@ def analyze_context(context, num_questions):
         )
 
 with gr.Blocks() as demo:
-    gr.Markdown("## Hệ thống sinh câu hỏi từ Context bằng mô hình Encoder-Decoder + spaCy NER")
+    gr.Markdown("## Hệ thống sinh câu hỏi")
 
     with gr.Row():
         with gr.Column(scale=4):
-            context_input = gr.Textbox(label="Nhập Context", lines=15, placeholder="Nhập đoạn văn bản...")
+            context_input = gr.Textbox(label="Nhập văn bản", lines=15, placeholder="Nhập đoạn văn bản...")
             elapsed_time_md = gr.Markdown(visible=False)
         with gr.Column(scale=1):
             model_choice = gr.Dropdown(