from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
import gradio as gr
 
# Tải mô hình và tokenizer
tokenizer = AutoTokenizer.from_pretrained("antphb/DS-Chatbox-gpt2-vietnamese-V3")
model = AutoModelForCausalLM.from_pretrained("antphb/DS-Chatbox-gpt2-vietnamese-V3")

# Tạo pipeline cho text-generation
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Định nghĩa hàm trả lời câu hỏi
def chat(input_text):
    response = pipe(input_text, max_length=150, num_return_sequences=1)
    return response[0]['generated_text']

iface = gr.Interface(fn=chat, 
                     inputs="text", 
                     outputs="text", 
                     title="Chat với GPT tiếng Việt", 
                     description="Nhập câu hỏi hoặc yêu cầu của bạn và mô hình sẽ trả lời",
                     live=True)


iface.launch()