import gradio as gr
from sentence_transformers import SentenceTransformer

# Tải mô hình
model = SentenceTransformer(model_name_or_path='Alibaba-NLP/gte-multilingual-base',
                            trust_remote_code=True)

def add_sentence(sentences, new_sentence):
    sentences.append(new_sentence)
    return sentences, "", f"Danh sách câu: {sentences}"

def gte_model(sentences):
    try:
        embeddings = model.encode(sentences)
        return embeddings.tolist()  # Chuyển numpy array sang danh sách
    except Exception as e:
        return f"Error: {str(e)}"

# Tạo giao diện Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Mô hình GTE Multilingual")
    gr.Markdown("Nhập từng câu, sau đó nhấn 'Thêm câu' để thêm vào danh sách. Nhấn 'Mã hóa' để nhận kết quả.")

    sentence_input = gr.Textbox(label="Nhập câu", placeholder="Nhập một câu tại đây...")
    add_button = gr.Button("Thêm câu")
    sentences_state = gr.State([])  # Lưu trữ danh sách các câu
    sentence_list_display = gr.Markdown("Danh sách câu: []")
    encode_button = gr.Button("Mã hóa")
    output = gr.JSON(label="Kết quả mã hóa")

    # Liên kết các sự kiện và hàm
    add_button.click(add_sentence, inputs=[sentences_state, sentence_input], outputs=[sentences_state, sentence_input, sentence_list_display])
    encode_button.click(gte_model, inputs=sentences_state, outputs=output)

# Khởi chạy giao diện
demo.launch()