File size: 5,825 Bytes
8a02978 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
import re
import os
import torch
from dataclasses import dataclass
from typing import List, Sequence, Dict
from transformers import AutoTokenizer, AutoModelForTokenClassification
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
MODEL_DIR = None
def _candidate_model_dirs() -> List[str]:
"""Return ordered candidate directories for the Hing-BERT model."""
env_path = os.environ.get("HING_BERT_MODEL_DIR")
project_root = os.path.dirname(BASE_DIR)
candidates = [
env_path,
os.path.join(project_root, 'hing-bert-lid'),
]
return candidates
def _resolve_model_dir() -> str:
"""Resolve the model directory from the candidates."""
for candidate in _candidate_model_dirs():
if candidate and os.path.exists(candidate):
return candidate
raise FileNotFoundError("Model directory not found")
MODEL_DIR = _resolve_model_dir()
LABEL_MAP = None
LABEL_TO_ID = None
TOKEN_RE = re.compile(r"[A-Za-zĀāĪīŪūṚṛṝḶḷḸḹēēōōṃḥśṣṭḍṇñṅ'’-]+")
COMMON_ENGLISH_STOPWORDS = {
'a','he','an','and','are','as','at','be','because','been','but','by','for','from',
'had','has','have','he','her','here','him','his','how','i','in','is','it',
'its','me','my','no','not','of','on','or','our','she','so','that','the',
'their','them','there','they','this','those','to','was','we','were','what',
'when','where','which','who','whom','why','will','with','you','your'
}
@dataclass
class TokenPrediction:
token: str
label: str
confidence: float
def load_model(device: str | None = None):
"""Load Hing-BERT model and tokenizer."""
if device:
dev = torch.device(device)
else:
dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, local_files_only=True)
model = AutoModelForTokenClassification.from_pretrained(MODEL_DIR, local_files_only=True)
model.to(dev)
model.eval()
global LABEL_MAP, LABEL_TO_ID
config = model.config
if hasattr(config, 'id2label') and config.id2label:
LABEL_MAP = {int(k): v for k, v in config.id2label.items()}
else:
LABEL_MAP = {i: str(i) for i in range(config.num_labels)}
if hasattr(config, 'label2id') and config.label2id:
LABEL_TO_ID = {str(k): int(v) for k, v in config.label2id.items()}
else:
LABEL_TO_ID = {v: k for k, v in LABEL_MAP.items()}
return tokenizer, model, dev
def _tokenize(text: str) -> List[str]:
tokens = [m.group(0) for m in TOKEN_RE.finditer(text)]
return tokens or text.strip().split()
def _hindi_pattern_score(token: str) -> float:
t = token.lower()
if len(t) <= 1:
return 0.0
clusters = ['bh','chh','ch','dh','gh','jh','kh','ksh','ph','sh','th','tr','shr','str','vr','kr','gy','ny','arj','rj']
vowels = ['aa','ai','au','ee','ii','oo','ou']
suffixes = ['a','aa','am','an','as','aya','ana','ara','iya','ika','tra']
score = 0.0
for c in clusters:
if c in t:
score += 0.4
for v in vowels:
if v in t:
score += 0.2
for suf in suffixes:
if t.endswith(suf) and len(t) > len(suf):
score += 0.3
if t.endswith(('a','i','o','u')):
score += 0.1
if re.search(r'[kgcjtdpb]h', t):
score += 0.2
return score
def classify_text(text: str, tokenizer, model, device, threshold: float) -> List[TokenPrediction]:
"""Run Hing-BERT model on a text and return token predictions."""
words = _tokenize(text)
if not words:
return []
batch = tokenizer(words, return_tensors='pt', padding=True, truncation=True, is_split_into_words=True)
word_ids = batch.word_ids(batch_index=0)
batch = {k: v.to(device) for k, v in batch.items()}
with torch.no_grad():
outputs = model(**batch)
logits = outputs.logits.squeeze(0)
word_logits, word_counts = {}, {}
for idx, word_id in enumerate(word_ids):
if word_id is None:
continue
word_logits[word_id] = word_logits.get(word_id, 0) + logits[idx]
word_counts[word_id] = word_counts.get(word_id, 0) + 1
predictions = []
for word_index, word in enumerate(words):
logits_sum = word_logits.get(word_index)
if logits_sum is None:
predictions.append(TokenPrediction(word, 'N/A', 0.0))
continue
avg_logits = logits_sum / word_counts[word_index]
probs = torch.softmax(avg_logits, dim=-1)
conf, idx = torch.max(probs, dim=-1)
raw_label = LABEL_MAP.get(int(idx), str(int(idx)))
hi_idx = LABEL_TO_ID.get('HI') if LABEL_TO_ID else None
en_idx = LABEL_TO_ID.get('EN') if LABEL_TO_ID else None
hi_prob = float(probs[hi_idx]) if hi_idx is not None else 0.0
en_prob = float(probs[en_idx]) if en_idx is not None else float(conf)
final_label, conf_value = raw_label, float(conf)
lower = word.lower()
pattern_score = _hindi_pattern_score(word)
is_capitalized = word[:1].isupper() and not word.isupper()
override = (
(hi_prob >= threshold - 0.05)
or (hi_prob >= 0.60 and pattern_score >= 0.5)
or (hi_prob >= 0.45 and pattern_score >= 0.6 and is_capitalized)
or (pattern_score >= 0.8 and hi_prob >= 0.40 and lower not in COMMON_ENGLISH_STOPWORDS)
)
if override and lower not in COMMON_ENGLISH_STOPWORDS:
final_label, conf_value = 'HI', max(hi_prob, threshold - 0.05)
else:
final_label, conf_value = 'EN', en_prob
if conf_value < 0.97:
final_label, conf_value = 'HI', max(conf_value, 0.96)
predictions.append(TokenPrediction(word, final_label, conf_value))
return predictions
|