embeddings / app.py
ЯRаСлав
Add JSON array batch mode support for embeddings
1b9b9da
Raw
History Blame Contribute Delete
2.13 kB
import gradio as gr
from sentence_transformers import SentenceTransformer
import json
# Високоякісна мультивалютна модель для 100+ мов
model_name = "intfloat/multilingual-e5-large"
model = SentenceTransformer(model_name)
def get_embeddings(input_data):
if not input_data or not input_data.strip():
return []
# Спробуємо розпарсити як JSON-масив для пакетної обробки (Batch Mode)
try:
data = json.loads(input_data.strip())
except Exception:
# Якщо не JSON, то це просто один рядок
data = input_data.strip()
if isinstance(data, list):
# Пакетний режим: додаємо префікс "query: ", якщо його немає
processed_texts = []
for text in data:
t = str(text).strip()
if not t.startswith("query:") and not t.startswith("passage:"):
t = f"query: {t}"
processed_texts.append(t)
# Повертає список векторів [[...], [...]]
return model.encode(processed_texts).tolist()
else:
# Режим одного рядка
t = str(data).strip()
if not t.startswith("query:") and not t.startswith("passage:"):
t = f"query: {t}"
# Повертає один вектор [...]
return model.encode(t).tolist()
# Gradio інтерфейс
demo = gr.Interface(
fn=get_embeddings,
inputs=gr.Textbox(
lines=5,
label="Текст або JSON-масив текстів",
placeholder='Введіть текст або ["Текст 1", "Текст 2"] для пакетної обробки...'
),
outputs=gr.JSON(label="Результат (Вектор або Масив векторів)"),
title="Premium Multilingual Text Embeddings API",
description=f"Model: {model_name} (1024 dimensions). Підтримує одиночні запити та JSON-масиви для пакетної обробки."
)
demo.launch()