geolip-bert-8192

from transformers import BertModel, BertTokenizer
from deep_bert_v3 import DeepBertV3, DeepBertV3Config
from safetensors.torch import load_file

config = DeepBertV3Config()
model = DeepBertV3(config)

# Load trained memory system weights
state = load_file("checkpoints/best/memory_system.safetensors")
model.load_state_dict(state, strict=False)

tokenizer = BertTokenizer.from_pretrained("google-bert/bert-large-uncased")
state = model.init_state(batch_size=1, device="cuda")

# Process document segment by segment
for segment_text in document_segments:
    tokens = tokenizer(segment_text, return_tensors="pt", max_length=480,
                       padding="max_length", truncation=True).to("cuda")
    outputs, state = model(tokens["input_ids"], tokens["attention_mask"], state)

# Final output encodes full document context
document_embedding = outputs["memory_output"]  # (1, 1024)