cognitor 0.2.0

Python SDK to extract relevant metrics from Small Language Model inference calls.

cognitor-py

cognitor-py is the Python SDK of our Cognitor platform. It is used to get detailed tracing and observability for Small Language Models applications. All metrics can be saved to a self-hosted instance of the Cognitor platform or to a local file.

At this time, cognitor-py supports HuggingFace transformers models.

Installation

pip install cognitor

Usage

Available logging modes

• log_type="database": Logs are sent to a self-hosted instance of the Cognitor platform. This requires a running instance of the Cognitor platform, otherwise the logs will fail to be sent.
• log_type="file": Logs are saved to a local file in JSON format. This is useful for development and debugging, but does not provide the same level of querying and visualization capabilities as the database mode

Inference monitoring

from cognitor import Cognitor

# Initialize your model and tokenizer
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
pipe = pipeline("text-generation", model=model_name, tokenizer=tokenizer)

cognitor = Cognitor(
    model_name=model_name,
    tokenizer=tokenizer,
    log_type="database", # or "file" for local file logging
    # Database connection parameters (only needed if log_type="database")
    host="localhost",
    port=5432,
    user="postgres",
    password="postgres",
    dbname="cognitor"
)

# Run inference within the monitor context
with cognitor.monitor() as m:
    input_text = "Once upon a time,"
    with m.track():
        output = pipe(input_text, max_length=50)
    m.capture(input_data=input_text, output=output)

Training monitoring

from cognitor import Cognitor, HFTrainingCallback
from transformers import Trainer, TrainingArguments

cognitor = Cognitor(
    model_name="my-model",
    log_type="database", # or "file" for local file logging
    # Database connection parameters (only needed if log_type="database")
    host="localhost",
    port=5432,
    user="postgres",
    password="postgres",
    dbname="cognitor",
)

training_args = TrainingArguments(
    output_dir="./output",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    eval_strategy="epoch",
    logging_strategy="steps",
    logging_steps=1,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    callbacks=[HFTrainingCallback(cognitor)],
)

# call new_training_run() to create a new training run ID for the callback-logged steps. If not called, all steps will be logged under the same training run ID.
cognitor.new_training_run()

trainer.train()