FULL CODE: https://github.com/frank-morales2020/AST/blob/main/TOPO_JEPA.ipynb
PAPER: https://zenodo.org/records/21365501
INFERENCE
# ============================================================================
# TOPO-JEPA INFERENCE - Load and Test Certified Model
# Sovereign Machine Lab | Frank Morales Aguilera, BEng, MEng, SMIEEE
# ============================================================================
# ============================================================================
# 1. IMPORTS
# ============================================================================
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import hf_hub_download
# ============================================================================
# 2. CONFIGURATION
# ============================================================================
HF_REPO_ID = 'frankmorales2020/topo-jepa-gpt-oss-20b-agnews'
BASE_MODEL_ID = 'openai/gpt-oss-20b'
HIDDEN_SIZE = 2880
JEPA_LATENT_DIM = 512
MAX_LENGTH = 64
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('=' * 75)
print('TOPO-JEPA INFERENCE')
print('=' * 75)
print(f'Device: {device}')
# ============================================================================
# 3. LOAD TOKENIZER
# ============================================================================
print('\n[1] Loading tokenizer...')
tokenizer = AutoTokenizer.from_pretrained(HF_REPO_ID, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
print('✅ Tokenizer loaded')
# ============================================================================
# 4. LOAD BACKBONE
# ============================================================================
print('\n[2] Loading GPT-OSS-20B backbone...')
base_model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL_ID,
trust_remote_code=True,
torch_dtype=torch.bfloat16
).to(device)
for param in base_model.parameters():
param.requires_grad = False
print('✅ Backbone loaded')
# ============================================================================
# 5. DEFINE MODEL ARCHITECTURE (Must match training)
# ============================================================================
class JEPAProjection(nn.Module):
def __init__(self, input_dim: int, latent_dim: int = JEPA_LATENT_DIM):
super().__init__()
self.projection = nn.Sequential(
nn.Linear(input_dim, latent_dim * 2),
nn.BatchNorm1d(latent_dim * 2),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(latent_dim * 2, latent_dim),
nn.BatchNorm1d(latent_dim)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return F.normalize(self.projection(x), dim=-1)
class JEPAPredictor(nn.Module):
def __init__(self, latent_dim: int = JEPA_LATENT_DIM):
super().__init__()
self.predictor = nn.Sequential(
nn.Linear(latent_dim, latent_dim * 2),
nn.BatchNorm1d(latent_dim * 2),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(latent_dim * 2, latent_dim),
nn.BatchNorm1d(latent_dim)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return F.normalize(self.predictor(x), dim=-1)
class TaskAwareModel(nn.Module):
def __init__(self, base_model: nn.Module, hidden_size: int = HIDDEN_SIZE):
super().__init__()
self.base_model = base_model
dev = next(base_model.parameters()).device
self.classifier_A = nn.Linear(hidden_size, 2, dtype=torch.bfloat16).to(dev)
self.classifier_B = nn.Linear(hidden_size, 2, dtype=torch.bfloat16).to(dev)
self.classifier_C = nn.Linear(hidden_size, 2, dtype=torch.bfloat16).to(dev)
self.online_projector = JEPAProjection(hidden_size, JEPA_LATENT_DIM).to(dev)
self.target_projector = JEPAProjection(hidden_size, JEPA_LATENT_DIM).to(dev)
self.predictor = JEPAPredictor(JEPA_LATENT_DIM).to(dev)
self.current_task = 'A'
def forward(self, input_ids, attention_mask=None):
outputs = self.base_model(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True
)
hidden_states = outputs.hidden_states[-1]
if attention_mask is not None:
seq_lens = torch.eq(attention_mask, 1).int().sum(-1) - 1
batch_idx = torch.arange(input_ids.shape[0], device=input_ids.device)
last_hidden = hidden_states[batch_idx, seq_lens, :]
else:
last_hidden = hidden_states[:, -1, :]
head = getattr(self, f'classifier_{self.current_task}')
return head(last_hidden)
def switch_task(self, task: str):
assert task in ('A', 'B', 'C')
self.current_task = task
# ============================================================================
# 6. LOAD CERTIFIED WEIGHTS
# ============================================================================
print('\n[3] Loading certified weights...')
model = TaskAwareModel(base_model=base_model)
# Download and load weights
weights_path = hf_hub_download(
repo_id=HF_REPO_ID,
filename='certified_topo_jepa_best.pt'
)
state_dict = torch.load(weights_path, map_location='cpu')
model.load_state_dict(state_dict, strict=False)
model.to(device)
model.eval()
print('✅ Certified weights loaded')
# ============================================================================
# 7. TASK LABELS
# ============================================================================
TASK_LABELS = {
'A': {0: 'World', 1: 'Sports'},
'B': {0: 'Business', 1: 'Sci/Tech'},
'C': {0: 'World', 1: 'Sci/Tech'}
}
# ============================================================================
# 8. INFERENCE FUNCTION
# ============================================================================
def predict(text: str, task: str = 'C', model=model, tokenizer=tokenizer, device=device):
"""
Run inference on a single text sample.
Args:
text: Input text string
task: Task to use ('A', 'B', or 'C')
model: The loaded model
tokenizer: The tokenizer
device: torch device
Returns:
dict: Prediction results with labels and confidence
"""
# Tokenize
inputs = tokenizer(
text,
max_length=MAX_LENGTH,
padding='max_length',
truncation=True,
return_tensors='pt'
).to(device)
# Switch task and predict
model.switch_task(task)
with torch.no_grad():
logits = model(input_ids=inputs.input_ids, attention_mask=inputs.attention_mask)
probs = F.softmax(logits.float(), dim=-1).squeeze().cpu().numpy()
pred_class = int(np.argmax(probs))
confidence = float(probs[pred_class])
label = TASK_LABELS[task][pred_class]
return {
'text': text,
'task': task,
'predicted_class': pred_class,
'label': label,
'confidence': confidence * 100,
'probabilities': {TASK_LABELS[task][i]: float(probs[i]) for i in range(2)}
}
# ============================================================================
# 9. TEST INFERENCE
# ============================================================================
print('\n' + '=' * 75)
print('RUNNING INFERENCE TESTS')
print('=' * 75)
# Test samples
test_samples = [
('A', 'The national team won the championship after a stunning comeback victory.'),
('B', 'Quarterly earnings beat analyst expectations driven by strong cloud revenue growth.'),
('C', 'New quantum computing startup secures massive initial funding round for enterprise deployment.'),
('C', 'Scientists discover new exoplanet in habitable zone using advanced telescope technology.'),
('C', 'Global trade negotiations face new challenges as emerging economies demand reforms.'),
]
print('\n[4] Running predictions...')
for task, text in test_samples:
result = predict(text, task)
print(f'\nTask {task} ({result["label"]}):')
print(f' Text: "{text[:60]}..."')
print(f' Confidence: {result["confidence"]:.2f}%')
print(f' Probabilities: {result["probabilities"]}')
# ============================================================================
# 10. BATCH INFERENCE FUNCTION
# ============================================================================
def predict_batch(texts: list, task: str = 'C'):
"""
Run inference on a batch of texts.
Args:
texts: List of text strings
task: Task to use ('A', 'B', or 'C')
Returns:
list: List of prediction results
"""
results = []
for text in texts:
results.append(predict(text, task))
return results
print('\n' + '=' * 75)
print('BATCH INFERENCE EXAMPLE')
print('=' * 75)
batch_texts = [
'The economy shows signs of recovery with strong job growth and consumer spending.',
'Breakthrough in renewable energy storage could revolutionize power grid infrastructure.',
'Political leaders gather for summit to discuss climate change and sustainable development.',
]
results = predict_batch(batch_texts, task='C')
for i, result in enumerate(results):
print(f'\nSample {i+1}: {result["label"]} ({result["confidence"]:.2f}%)')
print(f' "{result["text"][:50]}..."')
print('\n' + '=' * 75)
print('TOPO-JEPA INFERENCE COMPLETE')
print('=' * 75)
===========================================================================
TOPO-JEPA INFERENCE
===========================================================================
Device: cuda
[1] Loading tokenizer...
config.json: 100% 479/479 [00:00<00:00, 156kB/s][transformers] You are using a model of type `TOPO-JEPA` to instantiate a model of type ``. This may be expected if you are loading a checkpoint that shares a subset of the architecture (e.g., loading a `sam2_video` checkpoint into `Sam2Model`), but is otherwise not supported and can yield errors. Please verify that the checkpoint is compatible with the model you are instantiating.
tokenizer_config.json: 100% 377/377 [00:00<00:00, 139kB/s]tokenizer.json: 100% 27.9M/27.9M [00:02<00:00, 12.8MB/s]chat_template.jinja: 100% 16.7k/16.7k [00:00<00:00, 4.58MB/s]✅ Tokenizer loaded
[2] Loading GPT-OSS-20B backbone...
config.json: 100% 1.81k/1.81k [00:00<00:00, 590kB/s]model.safetensors.index.json: 100% 36.4k/36.4k [00:00<00:00, 11.3MB/s]Download complete: 100% 13.8G/13.8G [00:37<00:00, 248MB/s]Fetching 3 files: 100% 3/3 [00:37<00:00, 15.87s/it]Loading weights: 100% 411/411 [00:27<00:00, 11.25it/s]generation_config.json: 100% 177/177 [00:00<00:00, 54.2kB/s]✅ Backbone loaded
[3] Loading certified weights...
certified_topo_jepa_best.pt: 100% 41.9G/41.9G [02:02<00:00, 402MB/s]✅ Certified weights loaded
===========================================================================
RUNNING INFERENCE TESTS
===========================================================================
[4] Running predictions...
Task A (Sports):
Text: "The national team won the championship after a stunning come..."
Confidence: 97.70%
Probabilities: {'World': 0.022977370768785477, 'Sports': 0.977022647857666}
Task B (Sci/Tech):
Text: "Quarterly earnings beat analyst expectations driven by stron..."
Confidence: 97.81%
Probabilities: {'Business': 0.021948255598545074, 'Sci/Tech': 0.9780517220497131}
Task C (Sci/Tech):
Text: "New quantum computing startup secures massive initial fundin..."
Confidence: 97.00%
Probabilities: {'World': 0.029986508190631866, 'Sci/Tech': 0.9700134992599487}
Task C (Sci/Tech):
Text: "Scientists discover new exoplanet in habitable zone using ad..."
Confidence: 89.18%
Probabilities: {'World': 0.10818894952535629, 'Sci/Tech': 0.8918110132217407}
Task C (World):
Text: "Global trade negotiations face new challenges as emerging ec..."
Confidence: 95.63%
Probabilities: {'World': 0.956308901309967, 'Sci/Tech': 0.043691057711839676}
===========================================================================
BATCH INFERENCE EXAMPLE
===========================================================================
Sample 1: Sci/Tech (86.88%)
"The economy shows signs of recovery with strong jo..."
Sample 2: Sci/Tech (99.93%)
"Breakthrough in renewable energy storage could rev..."
Sample 3: World (99.93%)
"Political leaders gather for summit to discuss cli..."
===========================================================================
TOPO-JEPA INFERENCE COMPLETE
===========================================================================