Spaces:
Running
Running
File size: 1,660 Bytes
64428bf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
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()
|