Update app.py

app.py

@@ -4,17 +4,25 @@ from rdkit.Chem import Draw
 from rdkit import Chem
 import selfies as sf
 
+sf_input
 def greet1(name):
-   
-    return name
+   tokenizer = AutoTokenizer.from_pretrained("zjunlp/MolGen")
+    model = AutoModelForSeq2SeqLM.from_pretrained("zjunlp/MolGen")
     
-def greet2(name):
-    sf_output = ['[C][=C][C][=C][C][=C][Ring1][=Branch1]',
-    '[C][=C][C][=C][C][=C][C][=C][Ring1][=Branch1]',
-    '[C][=C][C][=C][C][=C][Ring1][=Branch1][C][=C][C][=C]',
-    '[C][=C][C][=C][C][=C][Ring1][=Branch1][C@H1][C][=C][C]',
-    '[C][=C][C][=C][C][=C][Ring1][=Branch1][C@H1][=C][C][=C]']
+    sf_input = tokenizer(name, return_tensors="pt")
+    
+    # beam search
+    molecules = model.generate(input_ids=sf_input["input_ids"],
+                              attention_mask=sf_input["attention_mask"],
+                              max_length=15,
+                              min_length=5,
+                              num_return_sequences=5,
+                              num_beams=5)
+
+    sf_output = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=True).replace(" ","") for g in molecules]
+    return sf_output
     
+def greet2(name):
     smis = [sf.decoder(i) for i in sf_output]
     
     mols = []
@@ -31,9 +39,6 @@ def greet2(name):
    
     return img
 
-def greet3(name):
-   
-    return name+"11"