コードでの外部LLMの作成¶

# endpoint = "https://app.datarobot.com/api/v2"
# token="<ADD_VALUE_HERE>"
# dr.Client(endpoint=endpoint, token=token)

dr.Client()

%%writefile custom_model/custom.py
import os
import pandas as pd


OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")


PROMPT_COLUMN_NAME = "promptText"
COMPLETION_COLUMN_NAME = "resultText"
ERROR_COLUMN_NAME = "error"


from openai import OpenAI
import pandas as pd


def load_model(*args, **kwargs):
    try:
        from datarobot_drum import RuntimeParameters
        api_key = RuntimeParameters.get("OPENAI_API_KEY")["apiToken"]
    except Exception:
        api_key = OPENAI_API_KEY
    return OpenAI(api_key=api_key)


def score(data, model, **kwargs):
    prompts = data["promptText"].tolist()
    responses = []
    errors = []
    for prompt in prompts:
        try: 
            response = model.chat.completions.create(
                model="gpt-4",
                messages=[
                    {"role": "user", "content": f"{prompt}"},
                ],
                temperature=0,
            )
            responses.append(response.choices[0].message.content)
            errors.append(None)
        except Exception as e:
            responses.append(None)
            error = f"{e.__class__.__name__}: {str(e)}"
            responses.append(None)
            errors.append(error)

    return pd.DataFrame({PROMPT_COLUMN_NAME: prompts, COMPLETION_COLUMN_NAME: responses, ERROR_COLUMN_NAME: errors})

test_data = pd.DataFrame(
        {PROMPT_COLUMN_NAME: ["What is the weather like where you live?"]}
    )
score(
    test_data, model=load_model()

)

%%writefile custom_model/requirements.txt
openai==1.35.4
datarobot-drum==1.11.5

%%writefile custom_model/model-metadata.yaml
---
name: OpenAI gpt-4
type: inference
targetType: textgeneration
runtimeParameterDefinitions:
  - fieldName: OPENAI_API_KEY
    type: credential
    credentialType: api_token
    description: Open API key

os.putenv("TARGET_NAME", COMPLETION_COLUMN_NAME)
!drum score --code-dir custom_model/ --target-type textgeneration --input custom_model/test_data.csv

# List all existing base environments
execution_environments = dr.ExecutionEnvironment.list()
execution_environments

for execution_environment in execution_environments:
    # print(execution_environment)
    if execution_environment.name == "[DataRobot] Python 3.11 GenAI":
        BASE_ENVIRONMENT = execution_environment
        environment_versions = dr.ExecutionEnvironmentVersion.list(
            execution_environment.id
        )
        break

BASE_ENVIRONMENT_VERSION = environment_versions[0]

print(BASE_ENVIRONMENT)
print(BASE_ENVIRONMENT_VERSION)
print(BASE_ENVIRONMENT.id)

CUSTOM_MODEL_NAME = "OpenAI Wrapper Model"
if CUSTOM_MODEL_NAME not in [c.name for c in dr.CustomInferenceModel.list()]:
    # Create a new custom model
    print("creating new custom model")
    custom_model = dr.CustomInferenceModel.create(
        name=CUSTOM_MODEL_NAME,
        target_type=dr.TARGET_TYPE.TEXT_GENERATION,
        target_name=COMPLETION_COLUMN_NAME,
        description="Wrapper for OpenAI completion",
        language="Python",
        is_training_data_for_versions_permanently_enabled=True,  # for latest updates to Model Registry in 9.0
    )
else:
    print("Custom MOdel Exists")
    custom_model = [c for c in dr.CustomInferenceModel.list() if c.name == CUSTOM_MODEL_NAME].pop()

# Create a new custom model version in DataRobot
print("Upload new version of model to DataRobot")
model_version = dr.CustomModelVersion.create_clean(
    custom_model_id=custom_model.id,
    base_environment_id=BASE_ENVIRONMENT.id,
    files=[
        "./custom_model/custom.py",
        "./custom_model/requirements.txt",
        "./custom_model/model-metadata.yaml",
        
    ],
    network_egress_policy=dr.NETWORK_EGRESS_POLICY.PUBLIC,
    runtime_parameter_values=[dr.models.custom_model_version.RuntimeParameterValue(field_name="OPENAI_API_KEY", type="credential", value="65baa1e82e6c8bb16561f72d")]
)

# Test a new version in DataRobot
print("Run test of new version in DataRobot")
custom_model_test = dr.CustomModelTest.create(
    custom_model_id=custom_model.id,
    custom_model_version_id=model_version.id,
    dataset_id=pred_test_dataset.id,
    max_wait=3600,  # 1 hour timeout
)
custom_model_test.overall_status

# Option 1
HOST = "https://app.datarobot.com"
for name, test in custom_model_test.detailed_status.items():
    print("Test: {}".format(name))
    print("Status: {}".format(test["status"]))
    print("Message: {}".format(test["message"]))

print(
    "Finished testing: "
    + HOST
    + "model-registry/custom-models/"
    + custom_model.id
    + "/assemble"
)

if CUSTOM_MODEL_NAME not in [m.name for m in dr.RegisteredModel.list()]:
    print("Creating New Registered Model")
    registered_model_version = dr.RegisteredModelVersion.create_for_custom_model_version(
        	model_version.id, name=CUSTOM_MODEL_NAME, description="LLM Wrapper Example from DataRobot Docs", 
            registered_model_name=CUSTOM_MODEL_NAME
    )
else:
    print("Using Existing Model")
    registered_model = [m for m in dr.RegisteredModel.list() if m.name == CUSTOM_MODEL_NAME].pop()
    registered_model_version = dr.RegisteredModelVersion.create_for_custom_model_version(
        	model_version.id, name=CUSTOM_MODEL_NAME, description="LLM Wrapper Example from DataRobot Docs", 
            registered_model_id=registered_model.id
    )

from datarobot_predict.deployment import predict

input_df =  pd.DataFrame({
        PROMPT_COLUMN_NAME: [
            "Give me some context on large language models and their applications?",
            "What is AutoML?"
        ],
    })


result_df, response_headers = predict(deployment, input_df)
result_df