Custom Model Integration

Custom Model Integration

Deeptrain is built to be model-agnostic. While the platform natively supports over 200 open-source and private models, you can extend its multi-modal capabilities to your proprietary inference endpoints or custom-trained models.

By integrating a custom model, you can leverage Deeptrain’s data connectors (Video, Audio, Flowcharts) while maintaining control over your specific inference logic and hosting environment.

Integration Overview

To integrate a custom model, you must define a provider configuration that points to your inference endpoint. Deeptrain expects a standard RESTful interface but allows for custom header injection and payload mapping to ensure compatibility with various serving frameworks (like FastAPI, TorchServe, or Hugging Face Inference Endpoints).

Configuring a Custom Provider

To connect a custom endpoint, initialize the model connector with a CustomProvider object. This object requires the base URL of your model and the specific authentication scheme your endpoint utilizes.

from deeptrain import ModelRegistry, CustomProvider

# Define your custom inference endpoint
my_custom_model = CustomProvider(
    name="my-proprietary-vision-llm",
    endpoint_url="https://api.yourdomain.ai/v1/inference",
    api_key="your_secure_token",
    model_type="vision-text" # Specifies the modality interface
)

# Register the model within the Deeptrain ecosystem
ModelRegistry.register(my_custom_model)

Request and Response Schema

To ensure seamless data flow, your custom endpoint should be configured to handle the multi-modal payloads generated by Deeptrain's processors.

Expected Request Format

Deeptrain sends a POST request with a JSON body. Depending on the modality (e.g., Image or Video), the data is typically passed as a Base64 encoded string or a URI.

{
  "model": "my-proprietary-vision-llm",
  "messages": [
    {
      "role": "user",
      "content": [
        { "type": "text", "text": "Describe this flowchart." },
        { "type": "image_url", "image_url": { "url": "data:image/jpeg;base64,..." } }
      ]
    }
  ],
  "stream": false
}

Expected Response Format

The integration layer expects a standard completion response. If your custom model uses a different schema, you can provide a response_mapping function:

def custom_parser(raw_response):
    # Map your custom API output to Deeptrain's standard format
    return {
        "text": raw_response["outputs"][0]["text"],
        "usage": raw_response["metadata"]["tokens"]
    }

ModelRegistry.register(my_custom_model, parser=custom_parser)

Supported Modalities for Custom Endpoints

When integrating a custom model, you can specify which of Deeptrain’s multi-modal processors the model should interface with:

• Text Integration: Standard LLM inference for handling context-window-extended data.
• Vision Integration: For custom models capable of processing images, flowcharts, or diagrams.
• Audio/Video Integration: Use Deeptrain’s Transcribe API to convert media into a format your custom model can ingest (e.g., text transcripts or frame-by-frame embeddings).

Testing the Connection

Once registered, you can verify the integration by running a diagnostic inference call through the Deeptrain client:

response = deeptrain.chat(
    model="my-proprietary-vision-llm",
    messages=[{"role": "user", "content": "Test connection"}]
)

if response.status_code == 200:
    print("Custom model successfully integrated.")

Security Considerations

• Header Injection: Use the headers parameter in CustomProvider to pass custom organization IDs or environment-specific metadata.
• Private VPCs: If your model is hosted within a private network, ensure the Deeptrain agent has the necessary outbound permissions to reach your internal load balancer.