Spaces:
Running on Zero
Running on Zero
| 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") | |
| 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() | |
| } | |
| 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() | |