Update prediction.py

prediction.py

@@ -38,9 +38,9 @@ scalers = {
 
 # Transformer model
 class TransformerRegressor(nn.Module):
-    def __init__(self, input_dim=2058, embedding_dim=768, ff_dim=1024, num_layers=2, output_dim=6):
+    def __init__(self, feat_dim=2058, embedding_dim=768, ff_dim=1024, num_layers=2, output_dim=6):
         super().__init__()
-        self.feat_proj = nn.Linear(input_dim, embedding_dim)
+        self.feat_proj = nn.Linear(feat_dim, embedding_dim)
         encoder_layer = nn.TransformerEncoderLayer(
             d_model=embedding_dim, nhead=8, dim_feedforward=ff_dim,
             dropout=0.1, batch_first=True
@@ -98,14 +98,13 @@ def get_morgan_fingerprint(smiles, radius=2, n_bits=2048):
     if mol is None:
         raise ValueError("Invalid SMILES string.")
     fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius, nBits=n_bits)
-    return np.array(fp, dtype=np.float32).reshape(-1,1)
+    return np.array(fp, dtype=np.float32).reshape(1,-1)
 
 # ChemBERTa embedding
-def get_chemberta_embedding(smiles: str, tokenizer, model):
+def get_chemberta_embedding(smiles: str, tokenizer, chemberta):
     inputs = tokenizer(smiles, return_tensors="pt", padding=True, truncation=True)
-    inputs = {k: v.to(device) for k, v in inputs.items()}
     with torch.no_grad():
-        outputs = model(**inputs)
+        outputs = chemberta(**inputs)
     return outputs.last_hidden_state.mean(dim=1).to(device)
 
 # Save to DB