Spaces:
Running
Running
File size: 2,257 Bytes
4d799f2 2af5b7d f63af71 4d799f2 2af5b7d 4d799f2 64428bf 4d799f2 64428bf 4d799f2 f63af71 64428bf 4d799f2 fe92162 4d799f2 f63af71 64428bf f63af71 |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import os
import sys
import tempfile
import pandas as pd
import gradio as gr
from smi_ted_light.load import load_smi_ted
# 1) Ajuste de paths para encontrar o inference
BASE_DIR = os.path.dirname(__file__)
INFERENCE_DIR = os.path.join(BASE_DIR, "smi-ted", "inference")
sys.path.append(INFERENCE_DIR)
# 2) Carregando o modelo
MODEL_DIR = os.path.join("smi-ted", "inference", "smi_ted_light")
model = load_smi_ted(
folder=MODEL_DIR,
ckpt_filename="smi-ted-Light_40.pt",
vocab_filename="bert_vocab_curated.txt",
)
# 3) Função única que gera embedding E CSV
def gerar_embedding_e_csv(smiles: str):
smiles = smiles.strip()
if not smiles:
# Se não digitou nada, retorna erro e esconde o botão de download
return {"erro": "digite uma sequência SMILES primeiro"}, gr.update(visible=False)
try:
# Gera o embedding
vetor = model.encode(smiles, return_torch=True)[0].tolist()
# Cria DataFrame e escreve CSV num arquivo temporário
df = pd.DataFrame([vetor])
tmp = tempfile.NamedTemporaryFile(suffix=".csv", delete=False)
df.to_csv(tmp.name, index=False)
tmp.close()
# Retorna: 1) JSON, 2) update para o File (path + visível)
return vetor, gr.update(value=tmp.name, visible=True)
except Exception as e:
# Em caso de erro interno, mostra mensagem e esconde o botão
return {"erro": str(e)}, gr.update(visible=False)
# 4) Montando a interface Blocks
with gr.Blocks() as demo:
gr.Markdown(
"""
## SMI-TED Embedding Generator
Cole uma sequência SMILES e receba:
1. O vetor embedding (768 floats) em JSON
2. Um botão para baixar esse vetor em CSV
"""
)
with gr.Row():
inp_smiles = gr.Textbox(label="SMILES", placeholder="Ex.: CCO")
btn = gr.Button("Gerar Embedding")
with gr.Row():
out_json = gr.JSON(label="Embedding (lista de floats)")
out_file = gr.File(label="Download do CSV", visible=False)
# 5) Ligando o botão à função única (dois outputs)
btn.click(
fn=gerar_embedding_e_csv,
inputs=inp_smiles,
outputs=[out_json, out_file]
)
if __name__ == "__main__":
demo.launch()
|