frankmorales2020/topological-ai-Kimi-VL-A3B-Thinking-multirun

CODE

https://github.com/frank-morales2020/AST/blob/main/KIMI2_TOPO.ipynb

INFERENCE

# ============================================================================
# INFERENCE CODE FOR TOPO-2026 CERTIFIED KIMI-VL-A3B-THINKING
# ============================================================================

import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
from huggingface_hub import hf_hub_download

# Configuration
REPO_ID = "frankmorales2020/topological-ai-Kimi-VL-A3B-Thinking-multirun"
BASE_MODEL_ID = "moonshotai/Kimi-VL-A3B-Thinking"
HIDDEN_SIZE = 2048
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class KimiLinear_InferenceModel(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.current_task = 'C' # Default to terminal task C

    def forward(self, input_ids, attention_mask=None):
        with torch.no_grad():
            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 set_task(self, task: str):
        assert task in ('A', 'B', 'C')
        self.current_task = task

# 1. Load Base Model and Tokenizer
print("[INFERENCE] Loading base model and tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.bfloat16,
).to(DEVICE)
base_model.eval()

# 2. Instantiate Task-Aware Wrapper
model = KimiLinear_InferenceModel(base_model=base_model).to(DEVICE)

# 3. Download and Load Certified Classifier Heads from Hugging Face Hub
print(f"[INFERENCE] Downloading certified heads from {REPO_ID}...")
weights_path = hf_hub_download(repo_id=REPO_ID, filename="classifier_heads.pt")
state_dict = torch.load(weights_path, map_location=DEVICE)
model.load_state_dict(state_dict)
model.eval()
print("[INFERENCE] Certified model loaded successfully!")

# 4. Test Inference Function
def predict_text(text: str, task: str = 'C') -> str:
    model.set_task(task)
    tokens = tokenizer(text, max_length=64, padding='max_length',
                       truncation=True, return_tensors='pt')
    input_ids = tokens.input_ids.to(DEVICE)
    attention_mask = tokens.attention_mask.to(DEVICE)
    
    with torch.no_grad():
        logits = model(input_ids=input_ids, attention_mask=attention_mask)
        pred_idx = torch.argmax(logits, dim=-1).item()
        
    # Task mapping labels
    labels_map = {
        'A': {0: "World", 1: "Sports"},
        'B': {0: "Business", 1: "Sci/Tech"},
        'C': {0: "World", 1: "Sci/Tech"}
    }
    return labels_map[task][pred_idx]

# Example test run
sample_text = "Scientists discover a new exoplanet orbiting a distant star system."
predicted_class = predict_text(sample_text, task='C')
print(f"\n[TEST] Sample Input: '{sample_text}'")
print(f"[TEST] Predicted Task C Class: {predicted_class}")

[INFERENCE] Loading base model and tokenizer...
Loading checkpoint shards: 100% 7/7 [00:06<00:00,  1.04s/it][INFERENCE] Downloading certified heads from frankmorales2020/topological-ai-Kimi-VL-A3B-Thinking-multirun...
classifier_heads.pt: 100% 32.8G/32.8G [02:20<00:00, 432MB/s][INFERENCE] Certified model loaded successfully!

[TEST] Sample Input: 'Scientists discover a new exoplanet orbiting a distant star system.'
[TEST] Predicted Task C Class: Sci/Tech