| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import AutoTokenizer, AutoModel |
| from datetime import datetime |
| import json |
| import os |
| import math |
|
|
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
|
|
| MODEL_DIR = 'model' |
| FULL_MODEL_PATH = os.path.join(MODEL_DIR, 'cascaded_best.pt') |
| CONFIG_PATH = os.path.join(MODEL_DIR, 'model_config.json') |
| TOKENIZER_PATH = os.path.join(MODEL_DIR, 'tokenizer') |
| BASE_MODEL_PATH = os.path.join(MODEL_DIR, 'base_model') |
|
|
| DICT_2 = os.path.join(MODEL_DIR, 'label2id_2.json') |
| DICT_4 = os.path.join(MODEL_DIR, 'label2id_4.json') |
| DICT_6 = os.path.join(MODEL_DIR, 'label2id_6.json') |
|
|
| RESULTS_PATH = os.path.join(MODEL_DIR, 'test_results.txt') |
|
|
|
|
| class ArcMarginProduct(nn.Module): |
| """ArcFace classifier (inference mode: no margin, just cosine * scale).""" |
| def __init__(self, in_features, out_features, s=30.0, m=0.30): |
| super().__init__() |
| self.s = s |
| self.m = m |
| self.weight = nn.Parameter(torch.FloatTensor(out_features, in_features)) |
| nn.init.xavier_uniform_(self.weight) |
| self.cos_m = math.cos(m) |
| self.sin_m = math.sin(m) |
| self.th = math.cos(math.pi - m) |
| self.mm = math.sin(math.pi - m) * m |
|
|
| def forward(self, x, label=None): |
| cosine = F.linear(F.normalize(x), F.normalize(self.weight)) |
| if label is not None and self.training: |
| sine = torch.sqrt(1.0 - cosine.pow(2).clamp(0, 1)) |
| phi = cosine * self.cos_m - sine * self.sin_m |
| phi = torch.where(cosine > self.th, phi, cosine - self.mm) |
| one_hot = torch.zeros_like(cosine) |
| one_hot.scatter_(1, label.view(-1, 1).long(), 1) |
| output = (one_hot * phi) + ((1.0 - one_hot) * cosine) |
| return output * self.s |
| return cosine * self.s |
|
|
|
|
| class CascadedClassifier(nn.Module): |
| """3-level cascaded classifier: 2 → 4 → 6 with ArcFace on level 6.""" |
| def __init__(self, base_model, hidden_size, n2, n4, n6, |
| dropout=0.15, arc_s=30.0, arc_m=0.3): |
| super().__init__() |
| self.base_model = base_model |
| self.drop = nn.Dropout(dropout) |
|
|
| self.head_2 = nn.Sequential( |
| nn.Linear(hidden_size, 256), nn.LayerNorm(256), nn.GELU(), |
| nn.Dropout(dropout), nn.Linear(256, n2)) |
|
|
| self.head_4_fusion = nn.Linear(hidden_size + n2, hidden_size) |
| self.head_4 = nn.Sequential( |
| nn.LayerNorm(hidden_size), nn.GELU(), nn.Dropout(dropout), |
| nn.Linear(hidden_size, 256), nn.GELU(), nn.Linear(256, n4)) |
|
|
| self.head_6_fusion = nn.Linear(hidden_size + n4, hidden_size) |
| self.head_6_feat = nn.Sequential( |
| nn.LayerNorm(hidden_size), nn.GELU(), nn.Dropout(dropout), |
| nn.Linear(hidden_size, 512), nn.GELU()) |
| self.head_6_arc = ArcMarginProduct(512, n6, s=arc_s, m=arc_m) |
|
|
| def forward(self, input_ids, attention_mask, label_6=None): |
| out = self.base_model(input_ids=input_ids, attention_mask=attention_mask) |
| cls_out = self.drop(out.last_hidden_state[:, 0, :]) |
|
|
| l2 = self.head_2(cls_out) |
| p2 = torch.softmax(l2, dim=1) |
| f4 = self.head_4_fusion(torch.cat([cls_out, p2], dim=1)) |
| l4 = self.head_4(f4) |
| p4 = torch.softmax(l4, dim=1) |
| f6 = self.head_6_fusion(torch.cat([cls_out, p4], dim=1)) |
| feat6 = self.head_6_feat(f6) |
| l6 = self.head_6_arc(feat6, label_6) |
| return l2, l4, l6 |
|
|
|
|
| def save_result(filepath, text, candidates, cascade_2, cascade_4): |
| """Append a single test result to the results txt file.""" |
| with open(filepath, 'a', encoding='utf-8') as f: |
| f.write(f"\n{'='*80}\n") |
| f.write(f"Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") |
| f.write(f"Input: {text}\n") |
| f.write(f"Cascade: {cascade_2} → {cascade_4}\n") |
| f.write(f"{'-'*80}\n") |
| f.write(f"{'#':<4} | {'Code':<12} | {'Score':<10} | {'P(6)':<8} | Chain\n") |
| f.write(f"{'-'*80}\n") |
| for i, c in enumerate(candidates[:5]): |
| cd = c['code'] |
| ch = f"{cd[:2]}({c['p2']:.2f})→{cd[:4]}({c['p4']:.2f})→{cd[:6]}({c['p6']:.2f})" |
| f.write(f"{i+1:<4} | {cd:<12} | {c['score']:.2e} | {c['p6']:.4f} | {ch}\n") |
| f.write(f"{'-'*80}\n") |
| if candidates[0]['score'] > 1e-3: |
| f.write("✅ Strong match.\n") |
| elif candidates[0]['p6'] < 0.1: |
| f.write("⚠️ Low confidence.\n") |
|
|
|
|
| def main(): |
| print("Loading bert-base-uncased FULL FT + ArcFace model (3-level, 6-digit)...") |
|
|
| if not os.path.exists(CONFIG_PATH): |
| print(f"Config not found: {CONFIG_PATH}. Train first.") |
| return |
|
|
| try: |
| config = json.load(open(CONFIG_PATH)) |
| model_name = config['model_name'] |
| hidden_size = config['hidden_size'] |
| max_seq_len = config['max_seq_len'] |
| counts = config['classes'] |
| dropout = config.get('dropout', 0.15) |
| arc_s = config.get('arcface_scale', 30.0) |
| arc_m = config.get('arcface_margin', 0.3) |
|
|
| l2id_2 = json.load(open(DICT_2)) |
| l2id_4 = json.load(open(DICT_4)) |
| l2id_6 = json.load(open(DICT_6)) |
|
|
| id2l_2 = {v: k for k, v in l2id_2.items()} |
| id2l_4 = {v: k for k, v in l2id_4.items()} |
| id2l_6 = {v: k for k, v in l2id_6.items()} |
|
|
| tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH) |
|
|
| if os.path.exists(BASE_MODEL_PATH): |
| base_model = AutoModel.from_pretrained(BASE_MODEL_PATH) |
| else: |
| base_model = AutoModel.from_pretrained(model_name) |
|
|
| model = CascadedClassifier( |
| base_model=base_model, hidden_size=hidden_size, |
| n2=counts['n2'], n4=counts['n4'], n6=counts['n6'], |
| dropout=dropout, arc_s=arc_s, arc_m=arc_m |
| ).to(device) |
|
|
| if os.path.exists(FULL_MODEL_PATH): |
| state_dict = torch.load(FULL_MODEL_PATH, map_location=device) |
| model.load_state_dict(state_dict, strict=False) |
|
|
| model.eval() |
| print(f"Loaded. Best val acc: {config.get('best_val_acc_6', 'N/A')}%") |
| print(f"Mode: {config.get('training_mode', 'N/A')}") |
|
|
| except Exception as e: |
| print(f"Error: {e}") |
| import traceback |
| traceback.print_exc() |
| return |
|
|
| |
| with open(RESULTS_PATH, 'a', encoding='utf-8') as f: |
| f.write(f"\n{'#'*80}\n") |
| f.write(f"Test session started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n") |
| f.write(f"Model: {config.get('model_name', 'N/A')}\n") |
| f.write(f"Architecture: {config.get('architecture', 'N/A')}\n") |
| f.write(f"Best val acc (6-digit): {config.get('best_val_acc_6', 'N/A')}%\n") |
| f.write(f"{'#'*80}\n") |
|
|
| print(f"\n📝 Results will be saved to: {RESULTS_PATH}") |
| print("\n--- HS Code Classification (3-level, 6-digit) ---") |
| print("Type description or 'q' to quit.\n") |
|
|
| while True: |
| try: |
| text = input("Description: ") |
| except (KeyboardInterrupt, EOFError): |
| break |
| if text.lower() in ('q', 'quit', 'exit') or not text.strip(): |
| if not text.strip(): |
| continue |
| break |
|
|
| enc = tokenizer(text, max_length=max_seq_len, padding='max_length', |
| truncation=True, return_tensors='pt') |
| ids = enc['input_ids'].to(device) |
| mask = enc['attention_mask'].to(device) |
|
|
| with torch.no_grad(): |
| with torch.amp.autocast('cuda'): |
| o2, o4, o6 = model(ids, mask) |
|
|
| p2 = F.softmax(o2, dim=1) |
| p4 = F.softmax(o4, dim=1) |
| p6 = F.softmax(o6, dim=1) |
|
|
| _, b2 = torch.max(p2, 1) |
| b2c = id2l_2.get(b2.item(), "") |
| _, b4 = torch.max(p4, 1) |
| b4c = id2l_4.get(b4.item(), "") |
|
|
| top_p, top_i = torch.topk(p6, 10, dim=1) |
|
|
| candidates = [] |
| for j in range(10): |
| idx = top_i[0][j].item() |
| prob6 = top_p[0][j].item() |
| code6 = id2l_6.get(idx, "Unk") |
|
|
| def get_prob(code_str, mapper, probs): |
| for k, v in mapper.items(): |
| if v == code_str: |
| return probs[0][k].item() |
| return 0.0 |
|
|
| pr2 = get_prob(code6[:2], id2l_2, p2) |
| pr4 = get_prob(code6[:4], id2l_4, p4) |
|
|
| eps = 1e-6 |
| score = (prob6**2) * ((pr4+eps)**0.5) * ((pr2+eps)**0.5) |
| if code6.startswith(b4c): |
| score *= 10.0 |
| elif code6[:2] == b2c: |
| score *= 5.0 |
|
|
| candidates.append({"code": code6, "score": score, "p6": prob6, |
| "p4": pr4, "p2": pr2}) |
|
|
| candidates.sort(key=lambda x: x["score"], reverse=True) |
|
|
| print(f"\n Cascade: {b2c} → {b4c}") |
| print("-" * 80) |
| print(f"{'#':<4} | {'Code':<12} | {'Score':<10} | {'P(6)':<8} | Chain") |
| print("-" * 80) |
| for i in range(min(5, len(candidates))): |
| c = candidates[i] |
| cd = c["code"] |
| ch = f"{cd[:2]}({c['p2']:.2f})→{cd[:4]}({c['p4']:.2f})→{cd[:6]}({c['p6']:.2f})" |
| print(f"{i+1:<4} | {cd:<12} | {c['score']:.2e} | {c['p6']:.4f} | {ch}") |
| print("-" * 80) |
|
|
| if candidates[0]['score'] > 1e-3: |
| print("✅ Strong match.") |
| elif candidates[0]['p6'] < 0.1: |
| print("⚠️ Low confidence.") |
|
|
| |
| save_result(RESULTS_PATH, text, candidates, b2c, b4c) |
| print(f" 📝 Saved to {RESULTS_PATH}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
