Skip to main content

Vision LLMs with LoRA fine-tuning.

Project description

PlimAI: Vision LLMs with Efficient LoRA Fine-Tuning

PyPI version Downloads License: MIT


PlimAI is a modular, research-friendly framework for building and fine-tuning Vision Large Language Models (LLMs) with efficient Low-Rank Adaptation (LoRA) support. It is designed for:

  • Researchers exploring new vision transformer architectures or fine-tuning strategies
  • Practitioners who want to adapt large vision models to custom datasets with limited compute
  • Developers looking for a clean, extensible codebase for vision-language AI

PlimAI provides a plug-and-play interface for LoRA, making it easy to experiment with parameter-efficient fine-tuning. The codebase is modular, so you can swap out or extend components like patch embedding, attention, or MLP heads.


🏗️ Architecture Overview

PlimAI is built around a modular Vision Transformer (ViT) backbone, with LoRA adapters injected into attention and MLP layers for efficient fine-tuning. The main components are:

graph TD
    A[Input Image] --> B[Patch Embedding]
    B --> C[+CLS Token & Positional Encoding]
    C --> D[Transformer Encoder]
    D --> E[LayerNorm]
    E --> F[MLP Head]
    F --> G[Output (e.g., Class logits)]
    subgraph LoRA Adapters
        D
    end

Main Modules

  • PatchEmbedding: Splits the image into patches and projects them into embedding space.
  • TransformerEncoder: Stack of transformer layers, each with multi-head self-attention and MLP blocks. LoRA adapters can be injected here.
  • LoRALinear: Low-rank adapters for efficient fine-tuning, only a small number of parameters are updated.
  • MLPHead: Final classification or regression head.
  • Config & Utils: Easy configuration and preprocessing utilities.

📦 Installation

pip install plimai

Or, for the latest version from source:

git clone https://github.com/plim-ai/plim.git
cd plim
pip install .

🧑‍💻 Quick Start

import torch
from plimai.models.vision_transformer import VisionTransformer
from plimai.utils.config import default_config

# Dummy image batch: batch_size=2, channels=3, height=224, width=224
x = torch.randn(2, 3, 224, 224)
model = VisionTransformer(
    img_size=default_config['img_size'],
    patch_size=default_config['patch_size'],
    in_chans=default_config['in_chans'],
    num_classes=default_config['num_classes'],
    embed_dim=default_config['embed_dim'],
    depth=default_config['depth'],
    num_heads=default_config['num_heads'],
    mlp_ratio=default_config['mlp_ratio'],
    lora_config=default_config['lora'],
)
out = model(x)
print('Output shape:', out.shape)

📚 Documentation


🧩 Module Breakdown

Module Description
PatchEmbedding Converts images to patch embeddings for transformer input
TransformerEncoder Stack of transformer layers with optional LoRA adapters
LoRALinear Low-rank adapters for parameter-efficient fine-tuning
MLPHead Output head for classification or regression
data.py Preprocessing and augmentation utilities
config.py Centralized configuration for model/training hyperparameters

🧪 Running Tests

pytest tests/

🤝 Contributing

We welcome contributions! Please see CONTRIBUTING.md for guidelines.

  • Open issues for bugs or feature requests
  • Submit pull requests for improvements
  • Star ⭐ the repo if you find it useful!

📄 License

This project is licensed under the MIT License. See the LICENSE file for details.


🌟 Acknowledgements

Directory Structure

plimai/
  models/
    vision_transformer.py
    lora.py
  components/
    patch_embedding.py
    attention.py
    mlp.py
  utils/
    data.py
    config.py
  example.py

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

plimai-0.1.5.tar.gz (7.5 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

plimai-0.1.5-py3-none-any.whl (8.6 kB view details)

Uploaded Python 3

File details

Details for the file plimai-0.1.5.tar.gz.

File metadata

  • Download URL: plimai-0.1.5.tar.gz
  • Upload date:
  • Size: 7.5 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.11

File hashes

Hashes for plimai-0.1.5.tar.gz
Algorithm Hash digest
SHA256 918bcf9cd8a0c7bd9002c4c72365d1570790591b8bc4295e256615bfbc2c9135
MD5 3428cd239641e4dc34a194af11a04c36
BLAKE2b-256 095f8a4de1f5c19d7a9b79196e72df43fb220f1bb6e2bfecadd13fb275fd3b0d

See more details on using hashes here.

File details

Details for the file plimai-0.1.5-py3-none-any.whl.

File metadata

  • Download URL: plimai-0.1.5-py3-none-any.whl
  • Upload date:
  • Size: 8.6 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.1.0 CPython/3.12.11

File hashes

Hashes for plimai-0.1.5-py3-none-any.whl
Algorithm Hash digest
SHA256 7ca85c0b8df26a673456ae23fd46704e95dd82365908cc0580cbfaf7b74fff27
MD5 1dbef377cacda53f3cc169df2a47403b
BLAKE2b-256 d97f5d5745039ac0430573bd87e5c8424e464e380c684d32ebf7e3ce63ec3298

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page