VietCat's picture
fix build error
9f79fe4
raw
history blame
1.58 kB
import os
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch
import gradio as gr
# Load model and tokenizer
model_id = "NlpHUST/gpt2-vietnamese"
try:
tokenizer = GPT2Tokenizer.from_pretrained(model_id)
model = GPT2LMHeadModel.from_pretrained(model_id)
except Exception as e:
print(f"Error loading model: {e}")
raise e
# Set device
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):
try:
inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
outputs = model.generate(
inputs,
max_length=max_length,
temperature=temperature,
do_sample=True,
num_beams=1
)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
except Exception as e:
return f"Error generating text: {e}"
# Gradio interface
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)