from mp_api.client import MPRester
from emmet.core.summary import HasProps
import openai
import langchain
from langchain import OpenAI
from langchain import agents
from langchain.agents import initialize_agent
from langchain.agents import Tool, tool
from langchain import LLMMathChain, SerpAPIWrapper
from gpt_index import GPTListIndex, GPTIndexMemory
from langchain import SerpAPIWrapper
from langchain.prompts.few_shot import FewShotPromptTemplate
from langchain.prompts.prompt import PromptTemplate
from langchain.vectorstores import FAISS, Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.prompts.example_selector import (MaxMarginalRelevanceExampleSelector,
SemanticSimilarityExampleSelector)
import requests
from rdkit import Chem
import pandas as pd
import os
class MAPITools:
def __init__(self):
self.model = 'text-ada-001' #maybe change to gpt-4 when ready
self.k=10
def get_material_atoms(self, formula):
'''Receives a material formula and returns the atoms symbols present in it separated by comma.'''
import re
pattern = re.compile(r"([A-Z][a-z]*)(\d*)")
matches = pattern.findall(formula)
atoms = []
for m in matches:
atom, count = m
count = int(count) if count else 1
atoms.append((atom, count))
return ",".join([a[0] for a in atoms])
def check_prop_by_formula(self, formula):
raise NotImplementedError('Should be implemented in children classes')
def search_similars_by_atom(self, atoms):
'''This function receives a string with the atoms separated by comma as input and returns a list of similar materials'''
atoms = atoms.replace(" ", "")
with MPRester(os.getenv("MAPI_API_KEY")) as mpr:
docs = mpr.summary.search(elements=atoms.split(','), fields=["formula_pretty", self.prop])
return docs
def create_context_prompt(self, formula):
raise NotImplementedError('Should be implemented in children classes')
def LLM_predict(self, prompt):
''' This function receives a prompt generate with context by the create_context_prompt tool and request a completion to a language model. Then returns the completion'''
llm = OpenAI(
model_name=self.model,
temperature=0.7,
n=1,
best_of=5,
top_p=1.0,
stop=["\n\n", "###", "#", "##"],
# model_kwargs=kwargs,
)
return llm.generate([prompt]).generations[0][0].text
def get_tools(self):
return [
Tool(
name = "Get atoms in material",
func = self.get_material_atoms,
description = (
"Receives a material formula and returns the atoms symbols present in it separated by comma."
)
),
Tool(
name = f"Checks if material is {self.prop_name} by formula",
func = self.check_prop_by_formula,
description = (
f"This functions searches in the material project's API for the formula and returns if it is {self.prop_name} or not."
)
),
# Tool(
# name = "Search similar materials by atom",
# func = self.search_similars_by_atom,
# description = (
# "This function receives a string with the atoms separated by comma as input and returns a list of similar materials."
# )
# ),
Tool(
name = f"Create {self.prop_name} context to LLM search",
func = self.create_context_prompt,
description = (
f"This function received a material formula as input and create a prompt to be inputed in the LLM_predict tool to predict if the material is {self.prop_name}."
if isinstance(self, MAPI_class_tools) else
f"This function received a material formula as input and create a prompt to be inputed in the LLM_predict tool to predict the {self.prop_name} of a material."
)
),
Tool(name = "LLM predictiom",
func = self.LLM_predict,
description = (
"This function receives a prompt generate with context by the create_context_prompt tool and request a completion to a language model. Then returns the completion"
)
)
]
class MAPI_class_tools(MAPITools):
def __init__(self, prop, prop_name, p_label, n_label):
super().__init__()
self.prop = prop
self.prop_name = prop_name
self.p_label = p_label
self.n_label = n_label
def check_prop_by_formula(self, formula):
f''' This functions searches in the material project's API for the formula and returns if it is {self.prop_name} or not'''
with MPRester(os.getenv("MAPI_API_KEY")) as mpr:
docs = mpr.summary.search(formula=formula, fields=["formula_pretty", self.prop])
if docs:
if docs[0].formula_pretty == formula:
return self.p_label if docs[0].dict()[self.prop] else self.n_label
return f"Could not find any material while searching {formula}"
def create_context_prompt(self, formula):
'''This function received a material formula as input and create a prompt to be inputed in the LLM_predict tool to predict if the formula is a stable material '''
elements = self.get_material_atoms(formula)
similars = self.search_similars_by_atom(elements)
similars = [
{'formula': ex.formula_pretty,
'prop': self.p_label if ex.dict()[self.prop] else self.n_label
} for ex in similars
]
examples = pd.DataFrame(similars).drop_duplicates().to_dict(orient="records")
example_selector = MaxMarginalRelevanceExampleSelector.from_examples(
examples,
OpenAIEmbeddings(),
FAISS,
k=self.k,
)
prefix=(
f'You are a bot who can predict if a material is {self.prop_name}.\n'
f'Given this list of known materials and the information if they are {self.p_label} or {self.n_label}, \n'
f'you need to answer the question if the last material is {self.prop_name}:'
)
prompt_template=PromptTemplate(
input_variables=["formula", "prop"],
template=f"Is {{formula}} a {self.prop_name} material?@@@\n{{prop}}###",
)
suffix = f"Is {{formula}} a {self.prop_name} material?@@@\n"
prompt = FewShotPromptTemplate(
# examples=examples,
example_prompt=prompt_template,
example_selector=example_selector,
prefix=prefix,
suffix=suffix,
input_variables=["formula"])
return prompt.format(formula=formula)
class MAPI_reg_tools(MAPITools):
# TODO: deal with units
def __init__(self, prop, prop_name):
super().__init__()
self.prop = prop
self.prop_name = prop_name
def check_prop_by_formula(self, formula):
''' This functions searches in the material project's API for the formula and returns if it is stable or not'''
with MPRester(os.getenv("MAPI_API_KEY")) as mpr:
docs = mpr.summary.search(formula=formula, fields=["formula_pretty", self.prop])
if docs:
if docs[0].formula_pretty == formula:
return docs[0].dict()[self.prop]
elif docs[0].dict()[self.prop] is None:
return f"There is no record of {self.prop_name} for {formula}"
return f"Could not find any material while searching {formula}"
def create_context_prompt(self, formula):
f'''This function received a material formula as input and create a prompt to be inputed in the LLM_predict tool to predict the {self.prop_name} of the material '''
elements = self.get_material_atoms(formula)
similars = self.search_similars_by_atom(elements)
similars = [
{'formula': ex.formula_pretty,
'prop': f"{ex.dict()[self.prop]:2f}" if ex.dict()[self.prop] is not None else None
} for ex in similars
]
examples = pd.DataFrame(similars).drop_duplicates().dropna().to_dict(orient="records")
example_selector = MaxMarginalRelevanceExampleSelector.from_examples(
examples,
OpenAIEmbeddings(),
FAISS,
k=self.k,
)
prefix=(
f'You are a bot who can predict the {self.prop_name} of a material .\n'
f'Given this list of known materials and the measurement of their {self.prop_name}, \n'
f'you need to answer the what is the {self.prop_name} of the material:'
'The answer should be numeric and finish with ###'
)
prompt_template=PromptTemplate(
input_variables=["formula", "prop"],
template=f"What is the {self.prop_name} for {{formula}}?@@@\n{{prop}}###",
)
suffix = f"What is the {self.prop_name} for {{formula}}?@@@\n"
prompt = FewShotPromptTemplate(
# examples=examples,
example_prompt=prompt_template,
example_selector=example_selector,
prefix=prefix,
suffix=suffix,
input_variables=["formula"])
return prompt.format(formula=formula)