import spaces
import gradio as gr
import torch
from sentence_transformers import SentenceTransformer

MODEL_ID = "AITeamVN/Vietnamese_Embedding_v2"

print(f"Loading {MODEL_ID} on ZeroGPU...")
model = SentenceTransformer(MODEL_ID)
model = model.to("cuda")

@spaces.GPU
def embed_text(text: str):
    if not text or not text.strip():
        return {"error": "Text is empty", "embedding": []}
    embedding = model.encode(text.strip(), normalize_embeddings=True)
    return {
        "model": MODEL_ID,
        "dim": int(len(embedding)),
        "embedding": embedding.tolist()
    }

@spaces.GPU
def embed_batch(texts_str: str):
    lines = [t.strip() for t in texts_str.split("\n") if t.strip()]
    if not lines:
        return {"error": "No texts provided", "embeddings": []}
    embeddings = model.encode(lines, normalize_embeddings=True)
    return {
        "model": MODEL_ID,
        "count": len(lines),
        "dim": int(embeddings.shape[1]),
        "embeddings": embeddings.tolist()
    }

with gr.Blocks(title="Vietnamese Embedding v2 API") as demo:
    gr.Markdown(f"# 🇻🇳 {MODEL_ID} Inference API")
    gr.Markdown("Dịch vụ tạo vector embedding 1024 chiều phục vụ RAG cho Chatbot HDSD.")
    
    with gr.Tab("Single Text"):
        text_input = gr.Textbox(label="Văn bản tiếng Việt", placeholder="Nhập câu cần tạo vector...", lines=2)
        btn_single = gr.Button("Trích xuất Vector", variant="primary")
        json_single = gr.JSON(label="Kết quả Vector")
        btn_single.click(fn=embed_text, inputs=text_input, outputs=json_single, api_name="embed")

    with gr.Tab("Batch Texts"):
        batch_input = gr.Textbox(label="Danh sách câu (mỗi dòng một câu)", lines=5)
        btn_batch = gr.Button("Trích xuất Batch Vector", variant="primary")
        json_batch = gr.JSON(label="Kết quả Danh sách Vector")
        btn_batch.click(fn=embed_batch, inputs=batch_input, outputs=json_batch, api_name="batch_embed")

demo.launch()