added inference + api wrapper

chemistral_api.py

@@ -0,0 +1,200 @@
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi import FastAPI, HTTPException, File, UploadFile, Form
+from fastapi.responses import JSONResponse, FileResponse
+from pydantic import BaseModel
+from typing import Optional
+import subprocess
+import os
+import logging
+from inference_transform import process_smiles, process_pdb, process_sdf, extract_and_convert_to_sdf, is_valid_smiles
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+app = FastAPI()
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=['*'],
+    allow_credentials=True,
+    allow_methods=['*'],
+    allow_headers=['*']
+)
+
+sdf_file_path = "/root/CHEMISTral7Bv0.3/example/Conformer3D_COMPOUND_CID_240.sdf"
+
+class InferenceRequest(BaseModel):
+    prompt: str
+    max_tokens: int = 256
+    temperature: float = 1.0
+
+@app.post("/predict_base")
+async def predict_base(
+    prompt: str = Form(...),
+    max_tokens: int = Form(256),
+    temperature: float = Form(1.0),
+    file: Optional[UploadFile] = File(None)
+):
+    try:
+        if file:
+            file_path = f"/tmp/{file.filename}"
+            with open(file_path, "wb") as f:
+                f.write(file.file.read())
+            if file.filename.endswith(".pdb"):
+                prompt += f" {process_pdb(file_path)}"
+            elif file.filename.endswith(".sdf"):
+                prompt += f" {process_sdf(file_path)}"
+        else:
+            try:
+                sdf_file = extract_and_convert_to_sdf(prompt)
+                if sdf_file:
+                    prompt += f" {sdf_file}"
+            except ValueError as e:
+                logger.info(str(e))
+
+        command = [
+            "python",
+            "/root/CHEMISTral7Bv0.3/mistral_chat_script.py",
+            "/root/mistral_models/7B-v0.3/",
+            prompt,
+            f"--max_tokens={max_tokens}",
+            f"--temperature={temperature}",
+            "--instruct"
+        ]
+
+        logger.info(f"Running command: {' '.join(command)}")
+        result = subprocess.run(command, capture_output=True, text=True)
+        
+        if result.returncode != 0:
+            logger.error(f"Command failed with return code {result.returncode}")
+            logger.error(f"stderr: {result.stderr}")
+            raise HTTPException(status_code=500, detail=result.stderr)
+        
+        response = result.stdout.strip()
+        sdf_file_path = "/root/CHEMISTral7Bv0.3/example/Conformer3D_COMPOUND_CID_240.sdf"
+
+        return {
+            "response": response,
+            "sdf_file_path": sdf_file_path
+        }
+    except Exception as e:
+        logger.exception("Exception occurred during inference.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.post("/predict")
+async def predict_alternative(
+    prompt: str = Form(...),
+    max_tokens: int = Form(256),
+    temperature: float = Form(1.0),
+    file: Optional[UploadFile] = File(None)
+):
+    try:
+        if file:
+            file_path = f"/tmp/{file.filename}"
+            with open(file_path, "wb") as f:
+                f.write(await file.read())
+            if file.filename.endswith(".pdb"):
+                prompt += f" {process_pdb(file_path)}"
+            elif file.filename.endswith(".sdf"):
+                prompt += f" {process_sdf(file_path)}"
+        else:
+            try:
+                sdf_file = extract_and_convert_to_sdf(prompt)
+                if sdf_file:
+                    prompt += f" {sdf_file}"
+            except ValueError as e:
+                logger.info(str(e))
+
+        command = [
+            "python",
+            "/root/CHEMISTral7Bv0.3/mistral_chat_script.py",
+            "/root/mistral_models/7B-v0.3/",
+            prompt,
+            f"--max_tokens={max_tokens}",
+            f"--temperature={temperature}",
+            "--instruct",
+            "--lora_path=/root/CHEMISTral7Bv0.3/runs/checkpoints/checkpoint_000300/consolidated/lora.safetensors"
+        ]
+        logger.info(f"Running command: {' '.join(command)}")
+        result = subprocess.run(command, capture_output=True, text=True)
+        if result.returncode != 0:
+            logger.error(f"Command failed with return code {result.returncode}")
+            logger.error(f"stderr: {result.stderr}")
+            raise HTTPException(status_code=500, detail=result.stderr)
+
+        response = result.stdout.strip()
+        sdf_file_path = "/root/CHEMISTral7Bv0.3/example/Conformer3D_COMPOUND_CID_240.sdf"
+
+        # Return the file as a direct download
+        return FileResponse(sdf_file_path, media_type='chemical/x-mdl-sdfile', filename="Conformer3D_COMPOUND_CID_240.sdf")
+
+    except Exception as e:
+        logger.exception("Exception occurred during inference.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+# @app.post("/predict")
+# async def predict_alternative(
+#     prompt: str = Form(...),
+#     max_tokens: int = Form(256),
+#     temperature: float = Form(1.0),
+#     file: Optional[UploadFile] = File(None)
+# ):
+#     try:
+#         global sdf_file_path
+#         if file:
+#             file_path = f"/tmp/{file.filename}"
+#             with open(file_path, "wb") as f:
+#                 f.write(file.file.read())
+#             if file.filename.endswith(".pdb"):
+#                 prompt += f" {process_pdb(file_path)}"
+#             elif file.filename.endswith(".sdf"):
+#                 prompt += f" {process_sdf(file_path)}"
+#         else:
+#             try:
+#                 sdf_file = extract_and_convert_to_sdf(prompt)
+#                 if sdf_file:
+#                     prompt += f" {sdf_file}"
+#             except ValueError as e:
+#                 logger.info(str(e))
+
+#         command = [
+#             "python",
+#             "/root/CHEMISTral7Bv0.3/mistral_chat_script.py",
+#             "/root/mistral_models/7B-v0.3/",
+#             prompt,
+#             f"--max_tokens={max_tokens}",
+#             f"--temperature={temperature}",
+#             "--instruct",
+#             "--lora_path=/root/CHEMISTral7Bv0.3/runs/checkpoints/checkpoint_000300/consolidated/lora.safetensors"
+#         ]
+
+#         logger.info(f"Running command: {' '.join(command)}")
+#         result = subprocess.run(command, capture_output=True, text=True)
+        
+#         if result.returncode != 0:
+#             logger.error(f"Command failed with return code {result.returncode}")
+#             logger.error(f"stderr: {result.stderr}")
+#             raise HTTPException(status_code=500, detail=result.stderr)
+        
+#         response = result.stdout.strip()
+#         sdf_file_path = "/root/CHEMISTral7Bv0.3/example/Conformer3D_COMPOUND_CID_240.sdf"
+
+#         return {
+#             "response": response,
+#             "sdf_file_path": sdf_file_path
+#         }
+#     except Exception as e:
+#         logger.exception("Exception occurred during inference.")
+#         raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/download_sdf")
+async def download_sdf():
+    try:
+        return FileResponse(path=sdf_file_path, filename="Conformer3D_COMPOUND_CID_240.sdf")
+    except Exception as e:
+        logger.exception("Exception occurred while sending SDF file.")
+        raise HTTPException(status_code=500, detail=str(e))
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

example/Conformer3D_COMPOUND_CID_240.sdf

@@ -0,0 +1,137 @@
+240
+  -OEChem-03012409263D
+
+ 14 14  0     0  0  0  0  0  0999 V2000
+    2.8466   -0.3870    0.0002 O   0  0  0  0  0  0  0  0  0  0  0  0
+    0.5644    0.2371    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -0.3437    1.2960    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
+    0.1013   -1.0787    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -1.7147    1.0393    0.0001 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -1.2698   -1.3354   -0.0001 C   0  0  0  0  0  0  0  0  0  0  0  0
+   -2.1777   -0.2764    0.0000 C   0  0  0  0  0  0  0  0  0  0  0  0
+    1.9937    0.5050   -0.0003 C   0  0  0  0  0  0  0  0  0  0  0  0
+    0.0016    2.3267    0.0000 H   0  0  0  0  0  0  0  0  0  0  0  0
+    0.7902   -1.9194   -0.0001 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -2.4218    1.8637    0.0001 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -1.6308   -2.3599   -0.0001 H   0  0  0  0  0  0  0  0  0  0  0  0
+   -3.2452   -0.4764    0.0000 H   0  0  0  0  0  0  0  0  0  0  0  0
+    2.2986    1.5653   -0.0006 H   0  0  0  0  0  0  0  0  0  0  0  0
+  1  8  2  0  0  0  0
+  2  3  2  0  0  0  0
+  2  4  1  0  0  0  0
+  2  8  1  0  0  0  0
+  3  5  1  0  0  0  0
+  3  9  1  0  0  0  0
+  4  6  2  0  0  0  0
+  4 10  1  0  0  0  0
+  5  7  2  0  0  0  0
+  5 11  1  0  0  0  0
+  6  7  1  0  0  0  0
+  6 12  1  0  0  0  0
+  7 13  1  0  0  0  0
+  8 14  1  0  0  0  0
+M  END
+> <PUBCHEM_COMPOUND_CID>
+240
+
+> <PUBCHEM_CONFORMER_RMSD>
+0.4
+
+> <PUBCHEM_CONFORMER_DIVERSEORDER>
+1
+
+> <PUBCHEM_MMFF94_PARTIAL_CHARGES>
+14
+1 -0.57
+10 0.15
+11 0.15
+12 0.15
+13 0.15
+14 0.06
+2 0.09
+3 -0.15
+4 -0.15
+5 -0.15
+6 -0.15
+7 -0.15
+8 0.42
+9 0.15
+
+> <PUBCHEM_EFFECTIVE_ROTOR_COUNT>
+1
+
+> <PUBCHEM_PHARMACOPHORE_FEATURES>
+2
+1 1 acceptor
+6 2 3 4 5 6 7 rings
+
+> <PUBCHEM_HEAVY_ATOM_COUNT>
+8
+
+> <PUBCHEM_ATOM_DEF_STEREO_COUNT>
+0
+
+> <PUBCHEM_ATOM_UDEF_STEREO_COUNT>
+0
+
+> <PUBCHEM_BOND_DEF_STEREO_COUNT>
+0
+
+> <PUBCHEM_BOND_UDEF_STEREO_COUNT>
+0
+
+> <PUBCHEM_ISOTOPIC_ATOM_COUNT>
+0
+
+> <PUBCHEM_COMPONENT_COUNT>
+1
+
+> <PUBCHEM_CACTVS_TAUTO_COUNT>
+1
+
+> <PUBCHEM_CONFORMER_ID>
+000000F000000001
+
+> <PUBCHEM_MMFF94_ENERGY>
+18.0728
+
+> <PUBCHEM_FEATURE_SELFOVERLAP>
+10.148
+
+> <PUBCHEM_SHAPE_FINGERPRINT>
+16714656 1 18409731763581766061
+18185500 45 18263078975662380655
+21040471 1 18338797793165636005
+23552423 10 18187929525178951646
+29004967 10 16200157598908099156
+369184 2 15195566792725449510
+5084963 1 18412544318583771616
+
+> <PUBCHEM_SHAPE_MULTIPOLES>
+158.77
+3.12
+1.41
+0.6
+1.61
+0.02
+0
+0.14
+0
+-0.45
+0
+-0.03
+-0.01
+0
+
+> <PUBCHEM_SHAPE_SELFOVERLAP>
+329.455
+
+> <PUBCHEM_SHAPE_VOLUME>
+90.5
+
+> <PUBCHEM_COORDINATE_TYPE>
+2
+5
+10
+
+$$$$

inference_transform.py

@@ -0,0 +1,48 @@
+import re
+from rdkit import Chem
+from rdkit.Chem import MolFromSmiles, SDWriter
+import logging
+from Bio import SeqIO
+
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+def process_smiles(smiles: str) -> str:
+    mol = MolFromSmiles(smiles)
+    if not mol:
+        raise ValueError(f"Invalid SMILES string: {smiles}")
+    
+    sdf_file = "/tmp/output.sdf"
+    writer = SDWriter(sdf_file)
+    writer.write(mol)
+    writer.close()
+    
+    return sdf_file
+
+def process_pdb(file_path: str) -> str:
+    sequences = []
+    with open(file_path, "r") as handle:
+        for record in SeqIO.parse(handle, "pdb-seqres"):
+            sequences.append(str(record.seq))
+    return " ".join(sequences)
+
+def process_sdf(file_path: str) -> str:
+    return file_path
+
+def extract_smiles(text: str) -> str:
+    smiles_pattern = r"([^J][0-9BCOHNSOPrIFla@+\-\[\]\(\)\\\/%=#$]{6,})"
+    matches = re.findall(smiles_pattern, text)
+    if matches:
+        return matches[0]
+    return ""
+
+def is_valid_smiles(smiles: str) -> bool:
+    mol = MolFromSmiles(smiles)
+    return mol is not None
+
+def extract_and_convert_to_sdf(text: str) -> str:
+    smiles = extract_smiles(text)
+    if smiles and is_valid_smiles(smiles):
+        return process_smiles(smiles)
+    raise ValueError("No valid SMILES string found in the text.")

mistral_chat_script.py

@@ -0,0 +1,38 @@
+import sys
+from pathlib import Path
+from mistral_inference.generate import generate
+from mistral_inference.model import Transformer
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+
+def run_chat(model_path: str, prompt: str, max_tokens: int = 256, temperature: float = 1.0, instruct: bool = True, lora_path: str = None):
+    # Find the correct tokenizer file
+    model_path = Path(model_path)
+    tokenizer_file = model_path / "tokenizer.model.v3"
+    
+    if not tokenizer_file.is_file():
+        raise FileNotFoundError(f"Tokenizer model file not found at {tokenizer_file}")
+
+    mistral_tokenizer = MistralTokenizer.from_file(str(tokenizer_file))
+    tokenizer = mistral_tokenizer.instruct_tokenizer.tokenizer
+
+    transformer = Transformer.from_folder(
+        model_path, max_batch_size=3, num_pipeline_ranks=1
+    )
+
+    if lora_path is not None:
+        transformer.load_lora(Path(lora_path))
+
+    tokens = tokenizer.encode(prompt, bos=True, eos=False)
+    generated_tokens, _ = generate(
+        [tokens],
+        transformer,
+        max_tokens=max_tokens,
+        temperature=temperature,
+        eos_id=tokenizer.eos_id,
+    )
+    answer = tokenizer.decode(generated_tokens[0])
+    print(answer)
+
+if __name__ == "__main__":
+    import fire
+    fire.Fire(run_chat)