remyxai/integer-programming

1import math2import json3import regex4import inspect5import guidance6from ast import literal_eval7from transformers import Tool8from ortools.linear_solver import pywraplp9 10guidance.llm = guidance.llms.OpenAI("gpt-4")11 12structure_program = guidance(13'''14{{#user~}}15{{description}}16Help me extract args from the data blob to apply the following algorithm:17{{code}}18 19----20 21{{~#each examples}}22Data Blob: {{this.input}}23Result: {{this.output}}24---25{{~/each}}26 27Please help me extract the input values from a given data blob into a JSON.28Data Blob: {{data_blob}}29Result: 30{{~/user}}31 32{{#assistant~}}33{{gen 'output'}}34{{~/assistant}}35''')36 37class IntegerProgrammingTool(Tool):38    name = "integer_programming_tool"39    description = """40    This tool solves an integer programming problem.41    Input is data_blob42    Output is the optimal solution as a string.43    """44    inputs = ["text"]45    outputs = ["text"]46    examples = [47        {'input': '''48        ,Space,Price49        Puzzle,1,250        BoardGame,5,1251 52        Constraint,10053        ''',54        'output': {55            "objective_coeffs": [56                [1, 2],57                [5, 15],58            ],59            "constraints": [100],60            "bounds": [61                [0, None],62                [0, None]63            ],64            "goal": "max"65          },66        },67        {'input': '''68        ,Space,Price69        Puzzle,3,1270        BoardGame,15,12071 72        Constraint,30073        ''',74        'output': {75            "objective_coeffs": [76                [3, 2],77                [15, 120],78            ],79            "constraints": [300],80            "bounds": [81                [0, None],82                [0, None]83            ],84            "goal": "max"85          },86        },87    ]88    def __call__(self, data_blob):89        code = inspect.getsourcelines(self.__call__)90        args = structure_program(91            description=self.description,92            code=code,93            examples=self.examples,94            data_blob=data_blob,95        )['output']96        pattern = regex.compile(r"\{(?:[^{}]|(?R))*\}")97        matches = pattern.findall(args)[0]98        args = literal_eval(matches)99        print(args)100        objective_coeffs = args['objective_coeffs']101        constraints = args['constraints']102        bounds = args['bounds']103        goal = args['goal']104 105        objective_coeffs = list(zip(*objective_coeffs))106        solver = pywraplp.Solver.CreateSolver("CBC")107        variables = [solver.IntVar(108        int(bounds[i][0]) if bounds[i][0] is not None else -math.inf,109        int(bounds[i][1]) if bounds[i][1] is not None else math.inf,110        f"x{i}") for i in range(len(objective_coeffs))]111 112        # Set objective function113        objective = solver.Objective()114        for i, coeff_list in enumerate(objective_coeffs):115            for j, coeff_value in enumerate(coeff_list):116                objective.SetCoefficient(variables[j], coeff_value)117        if goal == 'max':118            objective.SetMaximization()119        else:120            objective.SetMinimization()121 122        # Add constraints123        for i, constraint_value in enumerate(constraints):124            if constraint_value:125                constraint = solver.RowConstraint(0, constraint_value)126                for j, coeff in enumerate(objective_coeffs[i]):127                    constraint.SetCoefficient(variables[j], coeff)128        129        solver.Solve()130        131        solution = {"Objective value": objective.Value()}132        for i, variable in enumerate(variables):133            solution[f"x{i}"] = variable.solution_value()134        135        return json.dumps(solution)136