geolip-captionbert-8192

import torch
from transformers import AutoTokenizer
from caption_encoder import CaptionEncoder

# Load
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
model = CaptionEncoder(
    vocab_size=30522, max_len=8192, d_model=384,
    n_heads=6, n_layers=6, d_ff=1536, output_dim=768,
    dropout=0.0, pad_token_id=0)
model.load_state_dict(torch.load("best_model.pt", weights_only=True))
model.eval()

# Encode
texts = ["A cat sitting on a windowsill", "A dog playing fetch on the beach"]
tokens = tokenizer(texts, max_length=512, padding="max_length",
                   truncation=True, return_tensors="pt")
with torch.no_grad():
    embeddings = model(tokens["input_ids"], tokens["attention_mask"])

# embeddings: (2, 768) L2-normalized
similarity = embeddings[0] @ embeddings[1]
print(f"Similarity: {similarity:.3f}")