​
Register for an account

First, register for an account to get an API key. New accounts come with $5 to get started.

Once you’ve registered, go to settings and find your api key:

Watch our Full Demo Video You can watch our full demo here.

Google Colab example here.

Step 1. Install Flex AI on your OS:

pip install flex_ai

Step 2. Init your flex_ai instance

from flex_ai import FlexAI
client = FlexAI(api_key="xxxxxx")

​
Upload your first Dataset

In FlexAI, there are 3 types of Datasets:

  1. instruction
  2. chat
  3. text

You can read more about them in Datasets tutorial For now we will go with an instruction dataset. A datasets will have two .jsonl files, one for train and one for eval (eval is optional but recommended)

For example: create a train.jsonl file:

{"instruction": "Some important task with context ?", "output": "Boo" }
{"instruction": "Some important task with context ?", "output": "Boo" }
{"instruction": "Some important task with context ?", "output": "Boo" }

Now let’s create a new dataset

dataset = client.create_dataset("API Dataset New",
                               "instruction/train.jsonl",
                               "instruction/eval.jsonl")

You will get back a dataset id to use for fine tuning

​
Fine tune

You can view all our models in the Models Page You can then select a model and dataset and start training. You can also send fine tune directly from the Dashboard

task = client.create_finetune(name="My Task New",
                       dataset_id=dataset["id"],
                       model="meta-llama/Llama-3.2-1B-Instruct",
                       n_epochs=5,
                       train_with_lora=True,
                       lora_config={"lora_r": 64, "lora_alpha": 8, "lora_dropout": 0.1},
                       n_checkpoints_and_evaluations_per_epoch=1,
                       batch_size=4,
                       learning_rate=0.0001,
                       save_only_best_checkpoint=False
                      )

You can monitor the training progress and wait for completion:

Python
client.wait_for_task_completion(task_id=task["id"])

After the training is complete, you can retrieve the checkpoints:

Python
checkpoints = client.get_task_checkpoints(task_id=task["id"])

​
Create an OpenAI-compatible Endpoint

Once you have your checkpoints, you can create an endpoint to serve your fine-tuned model:

# Create a list of checkpoints to use
checkpoints_list = [{
    "id": checkpoint['id'],
    "name": "step_" + str(checkpoint['step'])
} for checkpoint in checkpoints]

# Create the endpoint
endpoint_id = client.create_multi_lora_endpoint(
    name="My Endpoint New",
    lora_checkpoints=checkpoints_list,
    compute="A100-40GB"
)

# Wait for the endpoint to be ready
endpoint = client.wait_for_endpoint_ready(endpoint_id=endpoint_id)

​
Use your Fine-tuned Model

You can use your fine-tuned model through the OpenAI-compatible API:

from openai import OpenAI

openai_endpoint_url = endpoint["url"]
client = OpenAI(
    api_key="your_api_key_here",
    base_url=f"{openai_endpoint_url}/v1"
)

completion = client.completions.create(
    model="meta-llama/Llama-3.2-1B-Instruct",
    prompt="Translate the following English text to French",
    max_tokens=60
)

print(completion.choices[0].text)

​
Full Example

Python
from flex_ai import FlexAI
from openai import OpenAI

# Initialize the Flex AI client
client = FlexAI(api_key="your_api_key_here")

# Create a new dataset
dataset = client.create_dataset(
    "API Dataset New",
    "instruction/train.jsonl",
    "instruction/eval.jsonl"
)

# Start a fine-tuning task
task = client.create_finetune(
    name="My Task New",
    dataset_id=dataset["id"],
    model="meta-llama/Llama-3.2-1B-Instruct",
    n_epochs=5,
    train_with_lora=True,
    lora_config={
        "lora_r": 64,
        "lora_alpha": 8,
        "lora_dropout": 0.1
    },
    n_checkpoints_and_evaluations_per_epoch=1,
    batch_size=4,
    learning_rate=0.0001,
    save_only_best_checkpoint=True
)

# Wait for training completion
client.wait_for_task_completion(task_id=task["id"])

# Get checkpoints
checkpoints = client.get_task_checkpoints(task_id=task["id"])
checkpoints_list = [{
    "id": checkpoint['id'],
    "name": "step_" + str(checkpoint['step'])
} for checkpoint in checkpoints]

# Create endpoint
endpoint_id = client.create_multi_lora_endpoint(
    name="My Endpoint New",
    lora_checkpoints=checkpoints_list,
    compute="A100-40GB"
)
endpoint = client.wait_for_endpoint_ready(endpoint_id=endpoint_id)

# Use the model
openai_client = OpenAI(
    api_key="your_api_key_here",
    base_url=f"{endpoint['url']}/v1"
)
completion = openai_client.completions.create(
    model="meta-llama/Llama-3.2-1B-Instruct",
    prompt="Translate the following English text to French",
    max_tokens=60
)

print(completion.choices[0].text)