| import os |
| import sys |
| import json |
| import time |
| import argparse |
| import numpy as np |
| import torch |
| from transformers import PreTrainedTokenizerFast |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from model import CompactGenderClassifier |
|
|
| SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) |
| PROJECT_ROOT = os.path.abspath(os.path.join(SCRIPT_DIR, "..")) |
|
|
| def find_first_existing(*paths): |
| for p in paths: |
| if os.path.exists(p): |
| return p |
| return paths[0] |
|
|
| CONFIG_PATH = find_first_existing( |
| os.path.join(SCRIPT_DIR, "config.json"), |
| os.path.join(SCRIPT_DIR, "config", "model_config.json"), |
| os.path.join(PROJECT_ROOT, "config", "model_config.json"), |
| os.path.join(PROJECT_ROOT, "config.json") |
| ) |
|
|
| TOKENIZER_DIR = find_first_existing( |
| os.path.join(SCRIPT_DIR, "tokenizer"), |
| os.path.join(PROJECT_ROOT, "models", "tokenizer"), |
| SCRIPT_DIR |
| ) |
|
|
| ONNX_PATH = find_first_existing( |
| os.path.join(SCRIPT_DIR, "model_int8.onnx"), |
| os.path.join(PROJECT_ROOT, "models", "model_int8.onnx"), |
| os.path.join(SCRIPT_DIR, "models", "model_int8.onnx") |
| ) |
|
|
| WEIGHTS_PATH = find_first_existing( |
| os.path.join(SCRIPT_DIR, "model_weights.pt"), |
| os.path.join(PROJECT_ROOT, "models", "model_weights.pt"), |
| os.path.join(SCRIPT_DIR, "models", "model_weights.pt") |
| ) |
|
|
| TITLE_PREFIXES = { |
| 'mr', 'mrs', 'ms', 'miss', 'dr', 'prof', 'sri', 'shri', 'smt', 'master', |
| 'md', 'mohd', 'sk', 'sekh', 'sheikh', 'syed', 'kazi', 'adv', 'er', 'ca' |
| } |
|
|
| class GenderClassifier: |
| """ |
| Dedicated binary Gender predictor (MALE vs FEMALE). |
| Supports ultra-fast ONNX INT8 CPU inference or PyTorch. |
| """ |
| def __init__(self, use_onnx=True): |
| if not os.path.exists(CONFIG_PATH): |
| raise FileNotFoundError(f"Configuration not found: {CONFIG_PATH}. Please train the model first.") |
|
|
| with open(CONFIG_PATH, "r", encoding="utf-8") as f: |
| self.cfg = json.load(f) |
|
|
| self.max_length = self.cfg.get("max_length", 20) |
| self.labels = {int(k): v for k, v in self.cfg.get("labels", {"0": "MALE", "1": "FEMALE"}).items()} |
|
|
| |
| tok_json = os.path.join(TOKENIZER_DIR, "tokenizer.json") |
| self.tokenizer = PreTrainedTokenizerFast( |
| tokenizer_file=tok_json, |
| do_lower_case=True, |
| pad_token="[PAD]", |
| unk_token="[UNK]", |
| cls_token="[CLS]", |
| sep_token="[SEP]", |
| mask_token="[MASK]" |
| ) |
|
|
| self.use_onnx = use_onnx and os.path.exists(ONNX_PATH) |
| if self.use_onnx: |
| import onnxruntime |
| session_opts = onnxruntime.SessionOptions() |
| session_opts.intra_op_num_threads = 1 |
| session_opts.inter_op_num_threads = 1 |
| session_opts.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL |
| self.session = onnxruntime.InferenceSession(ONNX_PATH, session_opts) |
| self.backend_name = "ONNX Runtime INT8 (Single-Thread)" |
| else: |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| self.model = CompactGenderClassifier( |
| vocab_size=self.cfg["vocab_size"], |
| emb_dim=self.cfg["emb_dim"], |
| hidden_dim=self.cfg["hidden_dim"], |
| num_layers=self.cfg["num_layers"], |
| num_heads=self.cfg["num_heads"], |
| max_len=self.max_length, |
| dropout=0.0 |
| ).to(self.device) |
| self.model.load_state_dict(torch.load(WEIGHTS_PATH, map_location=self.device, weights_only=True)) |
| self.model.eval() |
| self.backend_name = f"PyTorch ({self.device})" |
|
|
| def parse_name_tokens(self, text): |
| words = [w for w in str(text).strip().split() if len(w) > 0] |
| if not words: |
| return "", [] |
|
|
| idx = 0 |
| while idx < len(words): |
| cleaned = "".join([c for c in words[idx].lower() if c.isalpha()]) |
| if cleaned in TITLE_PREFIXES: |
| idx += 1 |
| else: |
| break |
|
|
| remaining = words[idx:] |
| if not remaining: |
| return words[0], [] |
| return remaining[0], remaining[1:] |
|
|
| def softmax(self, logits): |
| e_x = np.exp(logits - np.max(logits, axis=-1, keepdims=True)) |
| return e_x / np.sum(e_x, axis=-1, keepdims=True) |
|
|
| def predict(self, texts): |
| if isinstance(texts, str): |
| texts = [texts] |
|
|
| parsed_items = [self.parse_name_tokens(t) for t in texts] |
| query_names = [given if given else raw for raw, (given, _) in zip(texts, parsed_items)] |
|
|
| enc = self.tokenizer( |
| query_names, |
| padding="max_length", |
| max_length=self.max_length, |
| truncation=True, |
| return_tensors="np" |
| ) |
| input_ids = enc["input_ids"].astype(np.int32) |
| attention_mask = enc["attention_mask"].astype(np.int32) |
|
|
| t0 = time.time() |
| if self.use_onnx: |
| feeds = { |
| "input_ids": input_ids, |
| "attention_mask": attention_mask |
| } |
| logits = self.session.run(None, feeds)[0] |
| else: |
| with torch.no_grad(): |
| ids_t = torch.from_numpy(input_ids).to(self.device) |
| mask_t = torch.from_numpy(attention_mask).to(self.device) |
| logits = self.model(ids_t, mask_t).cpu().numpy() |
|
|
| elapsed_ms = (time.time() - t0) * 1000 / len(texts) |
| probs = self.softmax(logits) |
|
|
| results = [] |
| for i, text in enumerate(texts): |
| p = probs[i] |
| given_name, surnames = parsed_items[i] |
| pred_idx = int(np.argmax(p)) |
| label = self.labels.get(pred_idx, "UNKNOWN") |
| conf = float(p[pred_idx]) |
|
|
| results.append({ |
| "text": text, |
| "gender": label, |
| "confidence": conf, |
| "given_name": given_name, |
| "surnames": surnames, |
| "raw_probs": { |
| "male": float(p[0]), |
| "female": float(p[1]) |
| }, |
| "latency_ms": elapsed_ms |
| }) |
|
|
| return results |
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="CompactGenderClassifier Prediction Tool") |
| parser.add_argument("names", nargs="*", help="Name string(s) to classify as positional arguments") |
| parser.add_argument("--text", nargs="+", help="Name string(s) to classify") |
| parser.add_argument("-i", "--interactive", action="store_true", help="Start interactive CLI loop") |
| parser.add_argument("--use-pytorch", action="store_true", help="Force using PyTorch model instead of ONNX") |
| args = parser.parse_args() |
|
|
| clf = GenderClassifier(use_onnx=not args.use_pytorch) |
| print(f"--> Initialized {clf.backend_name}") |
|
|
| if args.interactive: |
| print("\n===============================================================") |
| print("CompactGenderClassifier Interactive Mode") |
| print("Type a name and press Enter. Type 'exit' or 'q' to quit.") |
| print("===============================================================\n") |
| while True: |
| try: |
| raw = input("Enter Name > ").strip() |
| if not raw: |
| continue |
| if raw.lower() in ["exit", "quit", "q"]: |
| print("Exiting interactive mode.") |
| break |
| r = clf.predict(raw)[0] |
| probs = r["raw_probs"] |
| given_str = f" [Given Name: {r['given_name']}]" if r["surnames"] else "" |
| print(f"--> Result: {r['gender']} ({r['confidence']*100:.1f}%) | Male: {probs['male']*100:.1f}% | Female: {probs['female']*100:.1f}%{given_str}\n") |
| except (KeyboardInterrupt, EOFError): |
| print("\nExiting interactive mode.") |
| break |
| return |
|
|
| input_names = args.names if args.names else args.text |
| if input_names: |
| names = input_names |
| else: |
| names = [ |
| "Supratim", |
| "Sayantika", |
| "Supratim Mondal", |
| "Sayantika Biswas", |
| "Alexander Hamilton", |
| "Priya Sharma", |
| "John Smith", |
| "Emily Watson", |
| "Dr. Rahul Verma", |
| "Ms. Fatima Sheikh" |
| ] |
|
|
| print("\n" + "=" * 95) |
| print(f"{'Input Name':<24} | {'Gender':<8} | {'Conf':<6} | {'Given Name':<14} | {'Probabilities'}") |
| print("=" * 95) |
|
|
| results = clf.predict(names) |
| for r in results: |
| probs = r['raw_probs'] |
| prob_str = f"Male: {probs['male']*100:.1f}% | Female: {probs['female']*100:.1f}%" |
| print(f"{r['text']:<24} | {r['gender']:<8} | {r['confidence']*100:5.1f}% | {r['given_name']:<14} | {prob_str}") |
|
|
| print("=" * 95) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
