Wip

bb4g/__init__.py

@@ -0,0 +1,161 @@
+import json
+import math
+from datetime import datetime
+from typing import List
+
+import requests
+from django.conf import settings as s
+
+from bridge.contracts import Classifier  # noqa: I001
+from bridge.dto import (
+    BuildResult,
+    EnviPyDTO,
+    EvaluationResult,
+    RunResult,
+    TransformationProductPrediction,
+)  # noqa: I001
+
+
+# Once stable these will be exposed by enviPy-plugins lib
+class BB4G(Classifier):
+    Config = None
+
+    def __init__(self, config=None):
+        super().__init__(config)
+        self.url = f"{s.BB4G_URL}"
+
+        self.token = self.acquire_token()
+        self.header = {
+            "Authorization": f"Bearer {self.token}",
+            "Content-Type": "application/json",
+        }
+        self.proxies = {
+            "http": s.HTTP_PROXY,
+            "https": s.HTTPS_PROXY,
+        }
+
+    def acquire_token(self):
+        BB4G_TENANT_ID = s.BB4G_TENANT_ID
+        BB4G_CLIENT_ID = s.BB4G_CLIENT_ID
+        BB4G_CLIENT_SECRET = s.BB4G_CLIENT_SECRET
+        BB4G_SCOPE = s.BB4G_SCOPE
+
+        BB4G_TOKEN_URL = f"https://login.microsoftonline.com/{BB4G_TENANT_ID}/oauth2/v2.0/token"
+
+        payload = {
+            "client_id": BB4G_CLIENT_ID,
+            "client_secret": BB4G_CLIENT_SECRET,
+            "scope": BB4G_SCOPE,
+            "grant_type": "client_credentials"
+        }
+
+        # No Proxy required, URL is whitelisted
+        res = requests.post(BB4G_TOKEN_URL, data=payload)
+
+        res.raise_for_status()
+
+        return res.json()["access_token"]
+
+    def start(self):
+        header = {
+            "Authorization": f"Bearer {self.token}",
+            "Content-Type": "application/json",
+        }
+
+        started = False
+        retries = 0
+        while not started and retries < 5:
+            res = requests.post(f"{self.url}/start", headers=header, data={}, proxies=self.proxies)
+
+            if res.status_code == 200:
+                started = True
+            elif res.status_code in [500, 502]:
+                retries += 1
+                import time
+                time.sleep(5)
+            else:
+                raise ValueError(f"Unexpected status code: {res.status_code}")
+
+    @classmethod
+    def requires_rule_packages(cls) -> bool:
+        return False
+
+    @classmethod
+    def requires_data_packages(cls) -> bool:
+        return False
+
+    @classmethod
+    def identifier(cls) -> str:
+        return "bb4g"
+
+    @classmethod
+    def name(cls) -> str:
+        return "BB4G Template Free Model"
+
+    @classmethod
+    def display(cls) -> str:
+        return "BB4G Template Free Model"
+
+    def build(self, eP: EnviPyDTO, *args, **kwargs) -> BuildResult | None:
+        return
+
+    def run(self, eP: EnviPyDTO, *args, **kwargs) -> RunResult:
+
+        # Ensure Service is running
+        self.start()
+
+        smiles = [c.smiles for c in eP.get_compounds()]
+        preds = self._post(smiles)
+
+        results = []
+
+        for substrate in preds.keys():
+            results.append(
+                TransformationProductPrediction(
+                    substrate=substrate,
+                    products=preds[substrate],
+                )
+            )
+
+        return RunResult(
+            producer=eP.get_context().url,
+            description=f"Generated at {datetime.now()}",
+            result=results,
+        )
+
+    def evaluate(self, eP: EnviPyDTO, *args, **kwargs) -> EvaluationResult:
+        pass
+
+    def build_and_evaluate(self, eP: EnviPyDTO, *args, **kwargs) -> EvaluationResult:
+        pass
+
+    def _post(self, smiles: List[str]) -> dict[str, dict[str, float]]:
+        header = {
+            "Authorization": f"Bearer {self.token}",
+            "Content-Type": "application/json",
+        }
+
+        result = {}
+
+        for smi in smiles:
+            data = {
+                "smiles": smi,
+                "sampling_alg": "exact",
+                "cutoff": -5,
+            }
+
+            resp = requests.post(f"{self.url}/compute", headers=header, data=json.dumps(data), proxies=self.proxies)
+
+            resp.raise_for_status()
+
+            for substrate, predictions in resp.json().items():
+                preds = {}
+
+                for pred in predictions:
+                    prod = pred["prediction"]
+                    prob = math.exp(pred["log_likelihood"])
+                    preds[prod] = prob
+
+                result[substrate] = preds
+
+        return result