Enzo Reis de Oliveira commited on
Commit
f63af71
·
1 Parent(s): 2af5b7d

Back withou csv

Browse files
Files changed (1) hide show
  1. app.py +28 -20
app.py CHANGED
@@ -1,46 +1,54 @@
1
- import os
2
- import sys
3
 
4
- # 1) Ajusta o path para o inference do SMI-TED **antes** de importar anything
5
  BASE_DIR = os.path.dirname(__file__)
6
  INFERENCE_DIR = os.path.join(BASE_DIR, "smi-ted", "inference")
7
- sys.path.append(INFERENCE_DIR)
8
 
9
- # Agora o python já sabe onde achar smi_ted_light
10
- import tempfile
11
- import pandas as pd
12
  import gradio as gr
13
- from smi_ted_light.load import load_smi_ted
 
14
 
15
  # 2) Caminho onde estão pesos e vocabulário
16
  MODEL_DIR = os.path.join("smi-ted", "inference", "smi_ted_light")
17
 
 
 
18
  model = load_smi_ted(
19
  folder=MODEL_DIR,
20
  ckpt_filename="smi-ted-Light_40.pt",
21
  vocab_filename="bert_vocab_curated.txt",
22
  )
23
 
 
24
  def gerar_embedding(smiles: str):
 
 
 
 
25
  smiles = smiles.strip()
26
  if not smiles:
27
- return {"erro": "digite uma sequência SMILES primeiro"}, None
 
28
  try:
29
- vetor = model.encode(smiles, return_torch=True)[0]
30
- emb = vetor.tolist()
31
- df = pd.DataFrame([emb])
32
- tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".csv", prefix="emb_")
33
- df.to_csv(tmp.name, index=False)
34
- tmp.close()
35
- return emb, tmp.name
36
  except Exception as e:
37
- return {"erro": str(e)}, None
38
 
 
 
39
  demo = gr.Interface(
40
  fn=gerar_embedding,
41
  inputs=gr.Textbox(label="SMILES", placeholder="Ex.: CCO"),
42
- outputs=[gr.JSON(), gr.File(label="Baixar CSV")],
43
- title="SMI-TED Embedding Generator",
 
 
 
 
44
  )
 
 
45
  if __name__ == "__main__":
46
- demo.launch(show_api=False)
 
1
+ import os, sys
 
2
 
 
3
  BASE_DIR = os.path.dirname(__file__)
4
  INFERENCE_DIR = os.path.join(BASE_DIR, "smi-ted", "inference")
5
+ sys.path.append(INFERENCE_DIR)
6
 
 
 
 
7
  import gradio as gr
8
+ from smi_ted_light.load import load_smi_ted
9
+
10
 
11
  # 2) Caminho onde estão pesos e vocabulário
12
  MODEL_DIR = os.path.join("smi-ted", "inference", "smi_ted_light")
13
 
14
+ # 3) Carrega o modelo SMI‑TED (Light)
15
+ # Se você renomeou o .pt ou o vocab, ajuste aqui.
16
  model = load_smi_ted(
17
  folder=MODEL_DIR,
18
  ckpt_filename="smi-ted-Light_40.pt",
19
  vocab_filename="bert_vocab_curated.txt",
20
  )
21
 
22
+ # 4) Função utilizada pela interface
23
  def gerar_embedding(smiles: str):
24
+ """
25
+ Recebe uma string SMILES e devolve o embedding (lista de 768 floats).
26
+ Em caso de erro, devolve um dicionário com a mensagem.
27
+ """
28
  smiles = smiles.strip()
29
  if not smiles:
30
+ return {"erro": "digite uma sequência SMILES primeiro"}
31
+
32
  try:
33
+ # model.encode devolve tensor shape (1, 768) quando return_torch=True
34
+ vetor_torch = model.encode(smiles, return_torch=True)[0]
35
+ return vetor_torch.tolist() # JSON‑serializável
 
 
 
 
36
  except Exception as e:
37
+ return {"erro": str(e)}
38
 
39
+
40
+ # 5) Define a interface Gradio
41
  demo = gr.Interface(
42
  fn=gerar_embedding,
43
  inputs=gr.Textbox(label="SMILES", placeholder="Ex.: CCO"),
44
+ outputs=gr.JSON(label="Embedding (lista de floats)"),
45
+ title="SMITED Embedding Generator",
46
+ description=(
47
+ "Cole uma sequência SMILES e receba o embedding gerado pelo modelo "
48
+ "SMI‑TED Light treinado pela IBM Research."
49
+ ),
50
  )
51
+
52
+ # 6) Roda localmente ou no Hugging Face Space
53
  if __name__ == "__main__":
54
+ demo.launch()