update

generate_poem.py

@@ -1,14 +1,76 @@
+import torch
 import gradio as gr
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 
+tokenizer = T5Tokenizer.from_pretrained("VietAI/vit5-base")
+model = T5ForConditionalGeneration.from_pretrained("Libosa2707/vietnamese-poem-t5")
 
-def generate_poem(text):
-    pass
+
+def generate_poem(input_text):
+    # Define the parameters for the generate function
+    min_length = 50
+    max_length = 100
+    rep_penalty = 1.2
+    temp = 0.7
+    top_k = 50
+    top_p = 0.92
+    no_repeat_ngram_size = 2
+
+    # Tokenize the input
+    input_ids = tokenizer(
+        input_text,
+        return_tensors="pt",
+        padding="max_length",
+        truncation=True,
+        max_length=42,
+    ).input_ids.to(model.device)
+
+    # Generate text
+    model.eval()
+    with torch.no_grad():
+        output = model.generate(
+            do_sample=True,
+            input_ids=input_ids,
+            min_length=min_length,
+            max_length=max_length,
+            top_p=top_p,
+            top_k=top_k,
+            temperature=temp,
+            repetition_penalty=rep_penalty,
+            no_repeat_ngram_size=no_repeat_ngram_size,
+            num_return_sequences=1,
+        )
+
+    # Process the generated text
+    gen = tokenizer.decode(
+        output[0], skip_special_tokens=False, clean_up_tokenization_spaces=False
+    )
+    sentences = gen.split("<unk>")
+    gen_poem = "\n".join(sentences).replace("<pad>", "").replace("</s>", "")
+    gen_poem = gen_poem.strip()
+
+    # Post-process the poem text
+    pretty_text = ""
+    for line in gen_poem.split("\n"):
+        line = line.strip()
+        if not line:
+            continue
+        line = line[0].upper() + line[1:]
+        pretty_text += line + "\n"
+
+    # Return the generated poem
+    return pretty_text
 
 
 generate_poem_interface = gr.Interface(
+    title="Làm thơ theo yêu cầu",
     fn=generate_poem,
     inputs=[
-        gr.components.Textbox(lines=1, placeholder="First words of the poem"),
+        gr.components.Textbox(
+            lines=1,
+            placeholder="Làm thơ với thể thơ tám chữ và tiêu đề mùa xuân nho nhỏ",
+            label="Yêu cầu về thể thơ và tiêu đề",
+        ),
     ],
     outputs="text",
 )