Using LLMs to Extract Structured Data: OpenAI Function Calling in Action

functions  = [
    {
        'name': 'MethodsList',
        'description': 'List of machine learning methods.',
        'parameters': {
            'type': 'object',
            'properties': {
                'methods': {
                    'title': 'Methods', 
                    'type': 'array', 
                    'items': {'type': 'string'},
            },
        }
        }, 'required': ['methodName', 'methodDescription']
    }
]

function_call = {"name": "MethodsList"}

system_prompt = """You are a machine learning engineer analyzing Kaggle competition solutions.
Your goal is to come up with the ultimate set of proven machine learning methods.
Review the winning solution description provided by user, and identify methods that worked.
Omit methods that didn't work. 
Don't add any methods not mentioned in the solution description.
Call the MethodsList function to save a list of methods that worked."""

MAX_TOKENS = 3500
encoding = tiktoken.encoding_for_model(GPT_MODEL)
def shorten(method_description):
    tokens = encoding.encode(method_description)
    return encoding.decode(tokens[:MAX_TOKENS])

from tenacity import retry, wait_random_exponential, stop_after_attempt
import json

@retry(wait=wait_random_exponential(min=1, max=40), stop=stop_after_attempt(3))
def chat_completion_request(system_prompt, method_description, functions=None, function_call=None, model=GPT_MODEL):
    short_prompt = shorten(method_description)
    messages = []
    messages.append({"role": "system", "content": system_prompt})
    messages.append({"role": "user", "content": short_prompt})
    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        functions=functions,
        function_call=function_call
    )
    try:
        res = response.choices[0].message.function_call.arguments
        data = json.loads(res)
        methods = data["methods"]
        return methods
    except:
        return response.choices[0].message

restructure_prompt = """You are a machine learning engineer analyzing Kaggle competition solutions. 
You are provided with a granular, unstructured list of machine learning methods that worked in various competitions.
Your goal is to restructure the list of methods to avoid duplicates or messy categories.
Do NOT use "Other" as a category. Always come up with a specific category. You can also decide to keep the original category as is.
Return a mapping of dirty method to cleaned up category, in csv format, for example:
XGBoost,XGBoost
xgb,XGBoost
Random Forest,Random Forest
CNN,CNNs
EfficientNet-B3,CNNs
SAINT,Transformers
DistillationLoss,DistillationLoss
BCE(smoothing=0.01),BCE
Here's the list of methods. Start your response with the mapping, don't add any extra description or text. Begin!
"""

from io import StringIO

completion = openai.ChatCompletion.create(
  model='gpt-4',
  messages=[
    {"role": "system", "content": restructure_prompt},
    {"role": "user", "content": methods_string},
  ]
)
csv_file = StringIO(completion.choices[0].message.content)

# Read the CSV-formatted string into a DataFrame
df = pd.read_csv(csv_file, header=None, names=['method', 'category'])

category_mapping = dict(zip(df.method, df.category))