update

app.py

@@ -1,66 +1,11 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from complete_poem import complete_poem_interface
+from generate_poem import generate_poem_interface
 
-tokenizer = AutoTokenizer.from_pretrained("vinai/phobert-base")
-# Define your models
-models = {
-    "Luc Bat": AutoModelForCausalLM.from_pretrained(
-        "Libosa2707/vietnamese-poem-luc-bat-gpt2"
-    ),
-    "Bay Chu": AutoModelForCausalLM.from_pretrained(
-        "Libosa2707/vietnamese-poem-bay-chu-gpt2"
-    ),
-    "Tam Chu": AutoModelForCausalLM.from_pretrained(
-        "Libosa2707/vietnamese-poem-tam-chu-gpt2"
-    ),
-    "Nam Chu": AutoModelForCausalLM.from_pretrained(
-        "Libosa2707/vietnamese-poem-nam-chu-gpt2"
-    ),
-}
-
-
-def generate_poem(text, style):
-    # Preprocess the input text
-    text = text.strip()
-    text = text.lower()
-
-    # Choose the model based on the selected style
-    model = models[style]
-
-    # Tokenize the input line
-    input_ids = tokenizer.encode(text, return_tensors="pt")[:, :-1]
-
-    # Generate text
-    output = model.generate(input_ids, max_length=100, do_sample=True, temperature=0.7)
-
-    # Decode the output
-    generated_text = tokenizer.decode(
-        output[:, input_ids.shape[-1] :][0], skip_special_tokens=True
-    )
-
-    text = text + " " + generated_text
-
-    # Post-process the output
-    text = text.replace("<unk>", "\n")
-    pretty_text = ""
-    for idx, line in enumerate(text.split("\n")):
-        line = line.strip()
-        if not line:
-            continue
-        line = line[0].upper() + line[1:]
-        pretty_text += line + "\n"
-
-    return pretty_text
-
-
-gradio_interface = gr.Interface(
-    fn=generate_poem,
-    inputs=[
-        gr.components.Textbox(lines=1, placeholder="First words of the poem"),
-        gr.components.Dropdown(
-            choices=["Luc Bat", "Bay Chu", "Tam Chu", "Nam Chu"], label="Style"
-        ),
-    ],
-    outputs="text",
+gr_interface = gr.TabbedInterface(
+    [complete_poem_interface, generate_poem_interface],
+    title="Thơ AI",
+    description="Ứng dụng mô hình ngôn ngữ cực lớn để sinh ra các bài thơ Việt Nam",
 )
-gradio_interface.launch()
+
+gr_interface.launch()

complete_poem.py

@@ -0,0 +1,65 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained("vinai/phobert-base")
+# Define your models
+models = {
+    "Luc Bat": AutoModelForCausalLM.from_pretrained(
+        "Libosa2707/vietnamese-poem-luc-bat-gpt2"
+    ),
+    "Bay Chu": AutoModelForCausalLM.from_pretrained(
+        "Libosa2707/vietnamese-poem-bay-chu-gpt2"
+    ),
+    "Tam Chu": AutoModelForCausalLM.from_pretrained(
+        "Libosa2707/vietnamese-poem-tam-chu-gpt2"
+    ),
+    "Nam Chu": AutoModelForCausalLM.from_pretrained(
+        "Libosa2707/vietnamese-poem-nam-chu-gpt2"
+    ),
+}
+
+
+def complete_poem(text, style):
+    # Preprocess the input text
+    text = text.strip()
+    text = text.lower()
+
+    # Choose the model based on the selected style
+    model = models[style]
+
+    # Tokenize the input line
+    input_ids = tokenizer.encode(text, return_tensors="pt")[:, :-1]
+
+    # Generate text
+    output = model.generate(input_ids, max_length=100, do_sample=True, temperature=0.7)
+
+    # Decode the output
+    generated_text = tokenizer.decode(
+        output[:, input_ids.shape[-1] :][0], skip_special_tokens=True
+    )
+
+    text = text + " " + generated_text
+
+    # Post-process the output
+    text = text.replace("<unk>", "\n")
+    pretty_text = ""
+    for idx, line in enumerate(text.split("\n")):
+        line = line.strip()
+        if not line:
+            continue
+        line = line[0].upper() + line[1:]
+        pretty_text += line + "\n"
+
+    return pretty_text
+
+
+complete_poem_interface = gr.Interface(
+    fn=complete_poem,
+    inputs=[
+        gr.components.Textbox(lines=1, placeholder="First words of the poem"),
+        gr.components.Dropdown(
+            choices=["Luc Bat", "Bay Chu", "Tam Chu", "Nam Chu"], label="Style"
+        ),
+    ],
+    outputs="text",
+)

generate_poem.py

@@ -0,0 +1,14 @@
+import gradio as gr
+
+
+def generate_poem(text):
+    pass
+
+
+generate_poem_interface = gr.Interface(
+    fn=generate_poem,
+    inputs=[
+        gr.components.Textbox(lines=1, placeholder="First words of the poem"),
+    ],
+    outputs="text",
+)