jina-v5-h256-lora-classification

Jina v5 H256 — LoRA Classification Adapters

Task-specific LoRA adapters for classification, built on top of the compressed gomyk/jina-v5-h256-distilled-conv embedding model.

Each MTEB Classification task has its own lightweight LoRA adapter (~405KB) + classification head, achieving +12.0%p average improvement over the frozen baseline embeddings.

LoRA Specification

How LoRA Works in This Model

For each attention projection (q/k/v/o) in each of the 6 layers:

  Original:  y = W_frozen @ x              W: [256, 256]
  With LoRA: y = W_frozen @ x + (α/r) * (x @ A) @ B

  Where:
    A: [256, 8]   — down-projection (trainable)
    B: [8, 256]   — up-projection (trainable)
    α/r = 16/8 = 2.0  — scaling factor

  Total per projection: 256×8 + 8×256 = 4,096 params
  Total per layer:      4 projections × 4,096 = 16,384 params
  Total (6 layers):     6 × 16,384 = 98,304 LoRA params
  + Classification head: Linear(256 → num_classes)

MTEB Classification Results

Average: 73.20% → 85.19% (+12.00%p)

Baseline: MTEB default evaluation (logistic regression on frozen embeddings) LoRA: MTEB evaluation with task-specific LoRA adapters applied to the embedding model

Per-Task Adapter Details

Training Configuration

Repository Structure

.
├── README.md
├── classification_results.json          # Summary of all results
├── AmazonCounterfactualClassification/
│   ├── lora_adapter.pt                  # LoRA A/B matrices
│   ├── classifier_head.pt              # Linear classification head
│   └── meta.json                        # Training metadata
├── Banking77Classification/
│   ├── ...
├── ImdbClassification/
│   ├── ...
├── MTOPDomainClassification/
│   ├── ...
├── MassiveIntentClassification/
│   ├── ...
├── MassiveScenarioClassification/
│   ├── ...
├── ToxicConversationsClassification/
│   ├── ...
└── TweetSentimentExtractionClassification/
    ├── ...

Usage

import torch
import torch.nn as nn
from transformers import AutoModel, AutoTokenizer

# 1. Base model load
base_model = AutoModel.from_pretrained("gomyk/jina-v5-h256-distilled-conv", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("gomyk/jina-v5-h256-distilled-conv", trust_remote_code=True)

# 2. LoRA adapter 적용
from huggingface_hub import hf_hub_download
import os

task = "Banking77Classification"  # 원하는 task 선택
repo_id = "gomyk/jina-v5-h256-lora-classification"

lora_path = hf_hub_download(repo_id, f"{task}/lora_adapter.pt")
head_path = hf_hub_download(repo_id, f"{task}/classifier_head.pt")
meta_path = hf_hub_download(repo_id, f"{task}/meta.json")

import json
with open(meta_path) as f:
    meta = json.load(f)

# 3. LoRA 부착
lora_state = torch.load(lora_path, map_location="cpu", weights_only=True)
idx = 0
for name, module in base_model.named_modules():
    for target in ["q_proj", "k_proj", "v_proj", "o_proj"]:
        child = getattr(module, target, None)
        if child is not None and isinstance(child, nn.Linear):
            A = lora_state[f"lora_{idx}_A"]
            B = lora_state[f"lora_{idx}_B"]
            scaling = lora_state[f"lora_{idx}_scaling"].item()
            # Merge into weight: W_new = W + scaling * (A @ B)^T
            child.weight.data += (scaling * (A @ B)).T
            idx += 1

# 4. Classification head
num_classes = meta["num_classes"]
classifier = nn.Sequential(nn.Dropout(0.1), nn.Linear(256, num_classes))
classifier.load_state_dict(torch.load(head_path, map_location="cpu", weights_only=True))

# 5. Inference
text = "I need to transfer money to another account"
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
with torch.no_grad():
    outputs = base_model(**inputs)
    hidden = outputs.last_hidden_state
    mask = inputs["attention_mask"].unsqueeze(-1).float()
    pooled = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)
    logits = classifier(pooled)
    pred = logits.argmax(dim=-1).item()
    print(f"Predicted class: {pred}")

License

Apache 2.0