BioMike commited on
Commit
6131fab
1 Parent(s): e65a0b2

Update utils/main_model.py

Browse files
Files changed (1) hide show
  1. utils/main_model.py +5 -2
utils/main_model.py CHANGED
@@ -8,8 +8,11 @@ class ChemicalConverter:
8
  self.mode = mode
9
  model_directory = os.path.abspath("models")
10
  model_path = os.path.join(model_directory, mode)
11
- if not os.path.exists(model_path):
12
- raise ValueError(f"Model path does not exist: {model_path}")
 
 
 
13
 
14
  self.model = MT5ForConditionalGeneration.from_pretrained(model_path)
15
  self.smiles_tokenizer = AutoTokenizer.from_pretrained("knowledgator/SMILES-FAST-TOKENIZER")
 
8
  self.mode = mode
9
  model_directory = os.path.abspath("models")
10
  model_path = os.path.join(model_directory, mode)
11
+
12
+ if mode == "SMILES2IUPAC":
13
+ model_path = "knowledgator/SMILES2IUPAC-canonical-base"
14
+ else:
15
+ model_path = "knowledgator/IUPAC2SMILES-canonical-small"
16
 
17
  self.model = MT5ForConditionalGeneration.from_pretrained(model_path)
18
  self.smiles_tokenizer = AutoTokenizer.from_pretrained("knowledgator/SMILES-FAST-TOKENIZER")