import os, sys

BASE_DIR = os.path.dirname(__file__)
INFERENCE_DIR = os.path.join(BASE_DIR, "smi-ted", "inference")
sys.path.append(INFERENCE_DIR) 

import gradio as gr
from smi_ted_light.load import load_smi_ted  


# 2) Caminho onde estão pesos e vocabulário
MODEL_DIR = os.path.join("smi-ted", "inference", "smi_ted_light")

# 3) Carrega o modelo SMI‑TED (Light)
#    Se você renomeou o .pt ou o vocab, ajuste aqui.
model = load_smi_ted(
    folder=MODEL_DIR,
    ckpt_filename="smi-ted-Light_40.pt",
    vocab_filename="bert_vocab_curated.txt",
)

# 4) Função utilizada pela interface
def gerar_embedding(smiles: str):
    """
    Recebe uma string SMILES e devolve o embedding (lista de 768 floats).
    Em caso de erro, devolve um dicionário com a mensagem.
    """
    smiles = smiles.strip()
    if not smiles:
        return {"erro": "digite uma sequência SMILES primeiro"}

    try:
        # model.encode devolve tensor shape (1, 768) quando return_torch=True
        vetor_torch = model.encode(smiles, return_torch=True)[0]
        return vetor_torch.tolist()  # JSON‑serializável
    except Exception as e:
        return {"erro": str(e)}


# 5) Define a interface Gradio
demo = gr.Interface(
    fn=gerar_embedding,
    inputs=gr.Textbox(label="SMILES", placeholder="Ex.: CCO"),
    outputs=gr.JSON(label="Embedding (lista de floats)"),
    title="SMI‑TED Embedding Generator",
    description=(
        "Cole uma sequência SMILES e receba o embedding gerado pelo modelo "
        "SMI‑TED Light treinado pela IBM Research."
    ),
)

# 6) Roda localmente ou no Hugging Face Space
if __name__ == "__main__":
    demo.launch()