# How to Set Classification and Summarization Models in Cognee

> Customize Cognee with your own classification and summarization models. Learn how to set models using cognee.config.set_classification_model() and cognee.config.set_summarization_model() easily.

- Repository: [Topoteretes/cognee](https://github.com/topoteretes/cognee)
- Tags: how-to-guide
- Published: 2026-03-16

---

**Use `cognee.config.set_classification_model()` and `cognee.config.set_summarization_model()` to replace the default Pydantic models that control how Cognee tags content and generates summaries.**

Cognee is an open-source knowledge graph engine that relies on configurable data models to classify and summarize content during the cognification pipeline. Learning how to set classification and summarization models enables you to customize Cognee's behavior without forking the repository or modifying environment variables. This guide explains the configuration architecture based on the actual source code in the `topoteretes/cognee` repository.

## Understanding the Default Models

Cognee's classification and summarization behavior is governed by two default Pydantic models defined in [`cognee/shared/data_models.py`](https://github.com/topoteretes/cognee/blob/main/cognee/shared/data_models.py):

- **DefaultContentPrediction** – A Pydantic model that represents a single class-label prediction for content classification
- **SummarizedContent** – A Pydantic model that holds a summary string and description for generated summaries

These models are referenced by the **CognifyConfig** class in [`cognee/modules/cognify/config.py`](https://github.com/topoteretes/cognee/blob/main/cognee/modules/cognify/config.py), which stores the active configuration in a singleton pattern. The runtime values determine how Cognee processes raw text into structured knowledge graph nodes.

## The Configuration API Pathway

Three files form the complete configuration pathway:

1. **[`cognee/shared/data_models.py`](https://github.com/topoteretes/cognee/blob/main/cognee/shared/data_models.py)** – Defines the base Pydantic models (`DefaultContentPrediction` and `SummarizedContent`)
2. **[`cognee/modules/cognify/config.py`](https://github.com/topoteretes/cognee/blob/main/cognee/modules/cognify/config.py)** – Contains the `CognifyConfig` class that holds the active model references
3. **[`cognee/api/v1/config/config.py`](https://github.com/topoteretes/cognee/blob/main/cognee/api/v1/config/config.py)** – Exposes the public configuration API with static setter methods

According to the source code, the setter methods are implemented at lines 64-73 in [`cognee/api/v1/config/config.py`](https://github.com/topoteretes/cognee/blob/main/cognee/api/v1/config/config.py), while a generic mapping table for string-based configuration exists at lines 28-33.

## Setting Custom Classification and Summarization Models

### Direct Setter Methods (Python API)

The most explicit way to configure models is through the dedicated setter methods. Your custom classes should inherit from the default models to ensure field compatibility.

```python
import cognee
from cognee.shared.data_models import DefaultContentPrediction, SummarizedContent

# Define custom models inheriting from defaults

class MyClassificationModel(DefaultContentPrediction):
    confidence_threshold: float = 0.8  # Add custom validation fields

class MySummarizationModel(SummarizedContent):
    metadata: dict = {}  # Extend with additional metadata

# Apply custom models at runtime

cognee.config.set_classification_model(MyClassificationModel)
cognee.config.set_summarization_model(MySummarizationModel)

# Verify the configuration

from cognee.modules.cognify.config import get_cognify_config
cfg = get_cognify_config()
print(cfg.classification_model)  # → <class '__main__.MyClassificationModel'>

print(cfg.summarization_model)   # → <class '__main__.MySummarizationModel'>

```

### Generic Configuration Setter

For dynamic configuration or when working with string-based settings, use the generic `set()` method:

```python
import cognee

cognee.config.set("classification_model", MyClassificationModel)
cognee.config.set("summarization_model", MySummarizationModel)

```

The `set()` method maps string keys to the appropriate dedicated setters via the configuration mapping table in [`cognee/api/v1/config/config.py`](https://github.com/topoteretes/cognee/blob/main/cognee/api/v1/config/config.py).

### Command Line Interface

You can also configure models via the CLI without writing Python code, assuming your custom models are importable:

```bash
cognee config set classification_model my_package.models.CustomClassifier
cognee config set summarization_model my_package.models.CustomSummarizer

```

The CLI command forwards the string path to `config.set()`, which resolves the class and updates the in-memory `CognifyConfig` singleton.

## Persisting Configuration Across Sessions

Because `CognifyConfig` exists only as an in-memory singleton, changes do not persist between Python process restarts. To ensure your custom models are always active, execute the configuration during application startup:

```python

# app/startup.py

import cognee
from my_models import CustomClassification, CustomSummarization

def initialize_cognee():
    cognee.config.set_classification_model(CustomClassification)
    cognee.config.set_summarization_model(CustomSummarization)
    # Proceed with cognify operations

```

Place this initialization in your application's entry point or [`__init__.py`](https://github.com/topoteretes/cognee/blob/main/__init__.py) to maintain consistent behavior across runs.

## Summary

- **Default models** (`DefaultContentPrediction` and `SummarizedContent`) are defined in [`cognee/shared/data_models.py`](https://github.com/topoteretes/cognee/blob/main/cognee/shared/data_models.py) and stored in `CognifyConfig`
- **Configuration API** methods live in [`cognee/api/v1/config/config.py`](https://github.com/topoteretes/cognee/blob/main/cognee/api/v1/config/config.py), providing both dedicated setters and a generic `set()` interface
- **Direct setters** (`set_classification_model` and `set_summarization_model`) offer type-safe model replacement at lines 64-73
- **CLI configuration** uses string key mapping defined at lines 28-33 to resolve model paths dynamically
- **In-memory state** requires initialization on every application startup for persistent behavior

## Frequently Asked Questions

### What fields must custom classification and summarization models implement?

Custom models must conform to the Pydantic `BaseModel` structure with fields matching `DefaultContentPrediction` (for classification) and `SummarizedContent` (for summarization) as defined in [`cognee/shared/data_models.py`](https://github.com/topoteretes/cognee/blob/main/cognee/shared/data_models.py). While inheriting from these classes ensures compatibility, any custom implementation must provide the same core attributes that Cognee's internal logic expects during the cognification process.

### Can I use different models for different datasets within the same application?

No, the current implementation uses a singleton `CognifyConfig` pattern where `get_cognify_config()` returns a single global instance. All classification and summarization operations within the same Python process will use the same model classes. To use different models for different datasets, you must reconfigure the models between operations or run separate processes.

### Where are the default model classes actually instantiated during processing?

The default models are referenced within the cognification pipeline defined in [`cognee/modules/cognify/config.py`](https://github.com/topoteretes/cognee/blob/main/cognee/modules/cognify/config.py). When content is processed, Cognee retrieves the active configuration via `get_cognify_config()` and uses the `classification_model` and `summarization_model` attributes stored in that singleton to validate and structure output data.

### Why does the CLI use string paths instead of Python imports?

The CLI accepts dot-notation string paths (e.g., `my_package.models.CustomModel`) because the generic `set()` method in [`cognee/api/v1/config/config.py`](https://github.com/topoteretes/cognee/blob/main/cognee/api/v1/config/config.py) dynamically imports and resolves these classes at runtime. This allows configuration without requiring users to write Python bootstrap code, though the referenced modules must be available in the Python path.