# Space Hugging Face — service d'embedding patembed pour ARIZ Copilot.
# Expose une API Gradio : embed(task, texts_json) -> vecteurs JSON.
# Le Space tourne sur CPU gratuit ; patembed-base (768 dim) suffit pour reranker ~100 abstracts.

import json

import gradio as gr
from sentence_transformers import SentenceTransformer

MODEL_ID = "datalyes/patembed-base"

# Préfixes d'instruction officiels du modèle (config_sentence_transformers.json, tâche problem2full)
PREFIXES = {
    "problem_query": "encode problem query for document retrieval: ",
    "document": "encode document for retrieval: ",
}

model = SentenceTransformer(MODEL_ID, device="cpu")


def embed(task: str, texts_json: str) -> str:
    """task: 'problem_query' | 'document' ; texts_json: JSON array de textes."""
    prefix = PREFIXES.get(task)
    if prefix is None:
        return json.dumps({"error": f"task inconnue: {task}"})
    texts = json.loads(texts_json)
    if not isinstance(texts, list) or len(texts) > 200:
        return json.dumps({"error": "texts doit etre un tableau JSON de 200 textes max"})
    vectors = model.encode(
        [prefix + t for t in texts],
        normalize_embeddings=True,
        batch_size=8,
    )
    return json.dumps({"embeddings": [v.tolist() for v in vectors]})


demo = gr.Interface(
    fn=embed,
    inputs=[
        gr.Dropdown(choices=list(PREFIXES.keys()), value="problem_query", label="Tâche"),
        gr.Textbox(lines=6, label="Textes (JSON array)"),
    ],
    outputs=gr.Textbox(label="Embeddings (JSON)"),
    title="patembed-base — embeddings brevets (ARIZ Copilot)",
    description=(
        "Service d'embedding basé sur datalyes/patembed-base (PatenTEB, Ayaou & Cavallucci 2025). "
        "Tâche problem2full : requête = énoncé de problème, documents = abstracts de brevets. "
        "Licence CC BY-NC-SA 4.0 — usage non commercial."
    ),
)

demo.launch()