VietCat's picture
fix request error
99c05e6
raw
history blame
1.3 kB
import os
os.environ["HF_HOME"] = "/tmp/hf_home"
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch
import gradio as gr
model_id = "NlpHUST/gpt2-vietnamese"
tokenizer = GPT2Tokenizer.from_pretrained(model_id)
model = GPT2LMHeadModel.from_pretrained(model_id)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()
def generate_text(prompt, max_length=100, temperature=1.0):
inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
outputs = model.generate(inputs, max_length=max_length, temperature=temperature, do_sample=True)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
demo = gr.Interface(
fn=generate_text,
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, 300, value=100, step=10, label="Độ dài tối đa"),
gr.Slider(0.5, 1.5, value=1.0, step=0.1, label="Nhiệt độ (Temperature)")
],
outputs="text",
title="Sinh văn bản tiếng Việt",
description="Dùng mô hình GPT-2 Vietnamese từ NlpHUST để sinh văn bản tiếng Việt.",
allow_flagging="never"
)
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0", server_port=7860)