|
|
import os |
|
|
import re |
|
|
import json |
|
|
from openai import OpenAI |
|
|
from tqdm import tqdm |
|
|
|
|
|
from rdkit import Chem |
|
|
from rdkit import RDLogger |
|
|
RDLogger.DisableLog('rdApp.*') |
|
|
|
|
|
def canonicalize_smiles(smiles): |
|
|
mol = Chem.MolFromSmiles(smiles) |
|
|
return Chem.MolToSmiles(mol, canonical=True) |
|
|
|
|
|
|
|
|
class Agent(): |
|
|
def __init__(self, url="http://localhost:8000/v1"): |
|
|
|
|
|
self.client = OpenAI( |
|
|
api_key="0", |
|
|
base_url=url |
|
|
) |
|
|
print ("[Client Init] done") |
|
|
|
|
|
def generate(self, query): |
|
|
|
|
|
messages = [ |
|
|
{ |
|
|
"role": "user", |
|
|
"content": query |
|
|
} |
|
|
] |
|
|
result = self.client.chat.completions.create( |
|
|
messages=messages, |
|
|
model="general") |
|
|
|
|
|
return result.choices[0].message.content |
|
|
|
|
|
class AgentMolLM(Agent): |
|
|
|
|
|
def __init__(self, url="http://localhost:8000/v1"): |
|
|
|
|
|
super(AgentMolLM, self).__init__(url=url) |
|
|
|
|
|
def generate(self, query, temperature=0.01, top_p=0.01, n_resp=1, add_no_think=False): |
|
|
|
|
|
query = query+" /no_think" if add_no_think else query |
|
|
messages = [ |
|
|
{ |
|
|
"role": "user", |
|
|
"content": query |
|
|
} |
|
|
] |
|
|
result = self.client.chat.completions.create( |
|
|
messages=messages, |
|
|
model="mol_lm", |
|
|
temperature=temperature, |
|
|
top_p=top_p, |
|
|
n=n_resp) |
|
|
|
|
|
if n_resp>1: |
|
|
results = list() |
|
|
for i in range(n_resp): |
|
|
results.append(result.choices[i].message.content) |
|
|
return results |
|
|
else: |
|
|
return result.choices[0].message.content |
|
|
|
|
|
|
|
|
class ExtractorSMILES(): |
|
|
def __init__(self): |
|
|
self.pattern = r"<SMILES>(.*?)</SMILES>" |
|
|
def extract(self, text): |
|
|
res = list() |
|
|
for smi in re.findall(self.pattern, text, re.DOTALL): |
|
|
res.append(smi.replace(";", ".")) |
|
|
return res |
|
|
|
|
|
class ExtractorFormula(): |
|
|
def __init__(self): |
|
|
self.pattern = r"<MOLFORMULA>(.*?)</MOLFORMULA>" |
|
|
def extract(self, text): |
|
|
return re.findall(self.pattern, text, re.DOTALL) |
|
|
|
|
|
class ExtractorIUPAC(): |
|
|
def __init__(self): |
|
|
self.pattern = r"<IUPAC>(.*?)</IUPAC>" |
|
|
def extract(self, text): |
|
|
return re.findall(self.pattern, text, re.DOTALL) |
|
|
|
|
|
class ExtractorValue(): |
|
|
def __init__(self): |
|
|
self.pattern = r"[+-]?\d*\.\d+" |
|
|
def extract(self, text): |
|
|
return re.findall(self.pattern, text) |
|
|
|
|
|
class ExtractorYN(): |
|
|
def __init__(self): |
|
|
pass |
|
|
|
|
|
def extract(self, text): |
|
|
if "yes" in text.lower(): |
|
|
return ["yes"] |
|
|
elif "no" in text.lower(): |
|
|
return ["no"] |
|
|
else: |
|
|
return ["N/A"] |
|
|
|
|
|
class ExtractorText(): |
|
|
def __init__(self): |
|
|
pass |
|
|
|
|
|
def extract(self, text): |
|
|
return [text] |
|
|
|
|
|
class Extractor(): |
|
|
def __init__(self, task_type="smiles"): |
|
|
if task_type == "smiles": |
|
|
self.extractor = ExtractorSMILES() |
|
|
elif task_type == "formula": |
|
|
self.extractor = ExtractorFormula() |
|
|
elif task_type == "iupac": |
|
|
self.extractor = ExtractorIUPAC() |
|
|
elif task_type == "value": |
|
|
self.extractor = ExtractorValue() |
|
|
elif task_type == "bool": |
|
|
self.extractor = ExtractorYN() |
|
|
else: |
|
|
self.extractor = ExtractorText() |
|
|
|
|
|
def extract(self, text): |
|
|
text = text.split("</think>\n")[-1] |
|
|
resp = self.extractor.extract(text=text) |
|
|
res = resp[0].strip() if len(resp) else "" |
|
|
return res |
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
agent = AgentMolLM() |
|
|
|
|
|
QUERY = "你是谁" |
|
|
ADD_NO_THINK = True |
|
|
|
|
|
|
|
|
resp = agent.generate(QUERY, add_no_think=ADD_NO_THINK, temperature=0.6, top_p=0.95, n_resp=2) |
|
|
|
|
|
print (resp) |
