Views
No views yet
answerdotai/ModernBERT-base backbone directly without fine-tuning. It's designed for zero-shot tag classification tasks where you want to use a pretrained embedding model for semantic similarity computation.answerdotai/ModernBERT-base1"""
2Example: Using ModernBERT-base for Tag Classification
3
4This example shows how to use the pretrained ModernBERT-base backbone
5for zero-shot tag classification using our module abstractions.
6
7Installation:
8 pip install git+https://github.com/Pieces/TAG-module.git@main
9 # Or: pip install -e .
10
11Note: ModernBERT requires Python < 3.14 due to torch.compile compatibility.
12"""
13
14import torch
15from tags_model.models.backbone import SharedTextBackbone
16from playground.validate_from_checkpoint import compute_ranked_tags
17
18# Load the pretrained backbone
19print("Loading ModernBERT-base...")
20backbone = SharedTextBackbone(
21 model_name="answerdotai/ModernBERT-base",
22 embedding_dim=768,
23 freeze_backbone=True,
24 pooling_mode="cls",
25 trust_remote_code=True, # Required for ModernBERT
26)
27backbone.eval()
28print("✓ Model loaded!")
29
30# Example query
31query_text = "Machine learning model for image classification using PyTorch"
32
33# Candidate tags to rank
34candidate_tags = [
35 "pytorch", "machine-learning", "deep-learning", "computer-vision",
36 "neural-networks", "cnn", "image-classification", "tensorflow",
37 "data-science", "python"
38]
39
40print(f"\nQuery: {query_text}")
41print(f"Candidate tags: {candidate_tags}\n")
42
43# Encode query and tags
44with torch.inference_mode():
45 query_emb = backbone.encode_texts([query_text], max_length=512, return_dict=False)[0]
46 tag_embs = backbone.encode_texts(candidate_tags, max_length=512, return_dict=False)
47
48print(f"Query embedding shape: {query_emb.shape}")
49print(f"Tag embeddings shape: {tag_embs.shape}")
50
51# Rank tags by similarity
52ranked_tags = compute_ranked_tags(
53 query_emb=query_emb,
54 pos_embs=torch.empty(0, 768), # No positives for zero-shot
55 neg_embs=torch.empty(0, 768), # No negatives for zero-shot
56 general_embs=tag_embs,
57 positive_tags=[],
58 negative_tags=[],
59 general_tags=candidate_tags,
60)
61
62# Display top-ranked tags
63print("\n" + "="*60)
64print("Top Ranked Tags:")
65print("="*60)
66for tag, rank, label, score in ranked_tags[:5]:
67 print(f"{rank:2d}. {tag:20s} (score: {score:.4f})")
68
69print("\n" + "="*60)
70print("Example complete!")
71
721# Install the repository first
2pip install git+https://github.com/Pieces/TAG-module.git@main
3# Or for local development:
4pip install -e .
5
6# Run the example
7python modernbert_example.py1@software{{tag_module,
2 title = {{TAG Module: Persona-Conditioned Contrastive Learning for Tag Classification}},
3 author = {{Your Name}},
4 year = {{2025}},
5 url = {{https://github.com/yourusername/tag-module}}
6}}