Metadata-Version: 2.4
Name: cognitor
Version: 0.2.0
Summary: Python SDK to extract relevant metrics from Small Language Model inference calls.
Author-email: Riccardo <riccardo@tanaos.com>
Project-URL: Homepage, https://github.com/riccardo/cognitor-py
Project-URL: Bug Tracker, https://github.com/riccardo/cognitor-py/issues
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3.10
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Classifier: Operating System :: OS Independent
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: transformers
Requires-Dist: psutil
Requires-Dist: torch
Requires-Dist: pydantic
Requires-Dist: psycopg2-binary
Requires-Dist: sqlalchemy
Requires-Dist: pydantic-settings>=2.8.1
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: black; extra == "dev"
Requires-Dist: isort; extra == "dev"
Requires-Dist: build; extra == "dev"
Requires-Dist: twine; extra == "dev"
Requires-Dist: python-dotenv>=1.0.1; extra == "dev"
Provides-Extra: evaluation
Requires-Dist: pytest; extra == "evaluation"
Requires-Dist: transformers; extra == "evaluation"
Requires-Dist: torch; extra == "evaluation"

# cognitor-py

`cognitor-py` is the Python SDK of our [Cognitor platform](https://github.com/tanaos/cognitor). 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](https://github.com/tanaos/cognitor) or to a local file.

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

## Installation

```bash
pip install cognitor
```

## Usage

### Available logging modes

- `log_type="database"`: Logs are sent to a self-hosted instance of the [Cognitor platform](https://github.com/tanaos/cognitor). 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

```python
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

```python
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()
```
