Views
No views yet
1
2# ============================================================================
3# INFERENCE CODE FOR TOPO-2026 CERTIFIED KIMI-VL-A3B-THINKING
4# ============================================================================
5
6import torch
7import torch.nn as nn
8from transformers import AutoModelForCausalLM, AutoTokenizer
9from huggingface_hub import hf_hub_download
10
11# Configuration
12REPO_ID = "frankmorales2020/topological-ai-Kimi-VL-A3B-Thinking-multirun"
13BASE_MODEL_ID = "moonshotai/Kimi-VL-A3B-Thinking"
14HIDDEN_SIZE = 2048
15DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
16
17class KimiLinear_InferenceModel(nn.Module):
18 def __init__(self, base_model: nn.Module, hidden_size: int = HIDDEN_SIZE):
19 super().__init__()
20 self.base_model = base_model
21 dev = next(base_model.parameters()).device
22 self.classifier_A = nn.Linear(hidden_size, 2, dtype=torch.bfloat16).to(dev)
23 self.classifier_B = nn.Linear(hidden_size, 2, dtype=torch.bfloat16).to(dev)
24 self.classifier_C = nn.Linear(hidden_size, 2, dtype=torch.bfloat16).to(dev)
25 self.current_task = 'C' # Default to terminal task C
26
27 def forward(self, input_ids, attention_mask=None):
28 with torch.no_grad():
29 outputs = self.base_model(
30 input_ids=input_ids,
31 attention_mask=attention_mask,
32 output_hidden_states=True
33 )
34 hidden_states = outputs.hidden_states[-1]
35 if attention_mask is not None:
36 seq_lens = torch.eq(attention_mask, 1).int().sum(-1) - 1
37 batch_idx = torch.arange(input_ids.shape[0], device=input_ids.device)
38 last_hidden = hidden_states[batch_idx, seq_lens, :]
39 else:
40 last_hidden = hidden_states[:, -1, :]
41 head = getattr(self, f'classifier_{self.current_task}')
42 return head(last_hidden)
43
44 def set_task(self, task: str):
45 assert task in ('A', 'B', 'C')
46 self.current_task = task
47
48# 1. Load Base Model and Tokenizer
49print("[INFERENCE] Loading base model and tokenizer...")
50tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True)
51if tokenizer.pad_token is None:
52 tokenizer.pad_token = tokenizer.eos_token
53
54base_model = AutoModelForCausalLM.from_pretrained(
55 BASE_MODEL_ID,
56 trust_remote_code=True,
57 torch_dtype=torch.bfloat16,
58).to(DEVICE)
59base_model.eval()
60
61# 2. Instantiate Task-Aware Wrapper
62model = KimiLinear_InferenceModel(base_model=base_model).to(DEVICE)
63
64# 3. Download and Load Certified Classifier Heads from Hugging Face Hub
65print(f"[INFERENCE] Downloading certified heads from {REPO_ID}...")
66weights_path = hf_hub_download(repo_id=REPO_ID, filename="classifier_heads.pt")
67state_dict = torch.load(weights_path, map_location=DEVICE)
68model.load_state_dict(state_dict)
69model.eval()
70print("[INFERENCE] Certified model loaded successfully!")
71
72# 4. Test Inference Function
73def predict_text(text: str, task: str = 'C') -> str:
74 model.set_task(task)
75 tokens = tokenizer(text, max_length=64, padding='max_length',
76 truncation=True, return_tensors='pt')
77 input_ids = tokens.input_ids.to(DEVICE)
78 attention_mask = tokens.attention_mask.to(DEVICE)
79
80 with torch.no_grad():
81 logits = model(input_ids=input_ids, attention_mask=attention_mask)
82 pred_idx = torch.argmax(logits, dim=-1).item()
83
84 # Task mapping labels
85 labels_map = {
86 'A': {0: "World", 1: "Sports"},
87 'B': {0: "Business", 1: "Sci/Tech"},
88 'C': {0: "World", 1: "Sci/Tech"}
89 }
90 return labels_map[task][pred_idx]
91
92# Example test run
93sample_text = "Scientists discover a new exoplanet orbiting a distant star system."
94predicted_class = predict_text(sample_text, task='C')
95print(f"\n[TEST] Sample Input: '{sample_text}'")
96print(f"[TEST] Predicted Task C Class: {predicted_class}")
97
981
2[INFERENCE] Loading base model and tokenizer...
3Loading 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...
4classifier_heads.pt: 100% 32.8G/32.8G [02:20<00:00, 432MB/s][INFERENCE] Certified model loaded successfully!
5
6[TEST] Sample Input: 'Scientists discover a new exoplanet orbiting a distant star system.'
7[TEST] Predicted Task C Class: Sci/Tech
8
9