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()