Logging ML Server Predictions

This code demonstrates FastAPI server with a single endpoint, assuming it serves for ML model predictions. The request to this endpoint automatically logged in Label Studio instance with await client.tasks.create() method.

Create a project in Label Studio

Let’s create a project in Label Studio to collect the predictions’ requests and responses from the ML model, and aim to task human reviewers to assess the correctness of the predictions.

1from label_studio_sdk.client import LabelStudio
2
3ls = LabelStudio(
4    base_url='http://localhost:8080',
5    api_key='YOUR-API-KEY'
6)
7
8project = ls.projects.create(
9    title='ML Observability Project',
10    label_config='''
11        <View>
12            <Text name="text" value="$text" />
13            <Text name="context" value="$context" />
14            <Text name="predictions" value="$predictions" />
15            <Choices name="correctness" toName="text">
16                <Choice value="Correct" />
17                <Choice value="Incorrect" />
18            <Choice value="Partially correct" />
19            </Choices>
20        </View>
21    '''
22)
23print(f'PROJECT ID: {project.id}')

Create FastAPI server

First install dependencies:

$pip install fastapi uvicorn

Let’s create a simple FastAPI server with a single endpoint /predict that accepts POST requests with JSON payload.

1import os
2from fastapi import FastAPI, HTTPException
3from pydantic import BaseModel
4from label_studio_sdk.client import AsyncLabelStudio
5
6app = FastAPI()
7
8# Initialize the async client with the API key and project ID from running Label Studio app
9# Remember to set LABEL_STUDIO_API_KEY and LABEL_STUDIO_PROJECT_ID environment variables
10client = AsyncLabelStudio(
11    base_url="http://localhost:8080",
12    api_key=os.getenv("LABEL_STUDIO_API_KEY"),
13)
14project_id = int(os.getenv("LABEL_STUDIO_PROJECT_ID"))
15
16
17# Some dummy input data
18class UserInput(BaseModel):
19    text: str
20    context: str
21
22
23@app.post("/predict")
24async def create_item(user_input: UserInput):
25    # Get model predictions
26    # Replace this with your model prediction code
27    # predictions = await model.predict(user_input.text, user_input.context)
28    predictions = '...'
29    data = {'text': user_input.text, 'context': user_input.context, 'predictions': predictions}
30    try:
31        task = await client.tasks.create(project=project_id, data=data)
32        return task
33    except Exception as e:
34        raise HTTPException(status_code=400, detail=str(e))

Run the server

$LABEL_STUDIO_API_KEY=your-api-key LABEL_STUDIO_PROJECT_ID=project-id uvicorn fastapi_server:app --reload

Logging predictions

Now you can send POST requests to http://localhost:8000/predict with JSON payload:

$curl -X POST "http://localhost:8000/predict" -H "Content-Type: application/json" -d '{"text": "example", "context": "context"}'

Open Label Studio project and see the tasks created from the server responses.