Spaces:
Running
Running
| """ | |
| Transformer-based Language Identification Model | |
| Uses Google MuRIL (Multilingual Universal Representations for Indian Languages) | |
| or XLM-RoBERTa for high-accuracy sequence classification. | |
| Requires GPU for practical training. Set enabled: true in config. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| from pathlib import Path | |
| from typing import Dict, List, Optional, Tuple | |
| import numpy as np | |
| import pandas as pd | |
| from loguru import logger | |
| try: | |
| import torch | |
| from torch.utils.data import Dataset, DataLoader | |
| from transformers import ( | |
| AutoTokenizer, | |
| AutoModelForSequenceClassification, | |
| TrainingArguments, | |
| Trainer, | |
| EarlyStoppingCallback, | |
| ) | |
| from transformers import DataCollatorWithPadding | |
| TRANSFORMERS_AVAILABLE = True | |
| except ImportError: | |
| TRANSFORMERS_AVAILABLE = False | |
| logger.warning("transformers/torch not installed. Transformer model unavailable.") | |
| from src.utils.languages import ( | |
| LANGUAGE_REGISTRY, | |
| CODE_TO_LABEL, | |
| LABEL_TO_CODE, | |
| NUM_CLASSES, | |
| ALL_LANGUAGE_CODES, | |
| ) | |
| # ───────────────────────────────────────────── | |
| # Dataset Wrapper | |
| # ───────────────────────────────────────────── | |
| class IndicLIDDataset: | |
| """Wraps a pandas DataFrame for HuggingFace Trainer.""" | |
| def __init__(self, df: pd.DataFrame, tokenizer, max_length: int = 128): | |
| self.texts = df["text"].tolist() | |
| self.labels = df["label"].tolist() | |
| self.tokenizer = tokenizer | |
| self.max_length = max_length | |
| def __len__(self): | |
| return len(self.texts) | |
| def __getitem__(self, idx): | |
| encoding = self.tokenizer( | |
| self.texts[idx], | |
| max_length=self.max_length, | |
| truncation=True, | |
| padding="max_length", | |
| return_tensors="pt", | |
| ) | |
| return { | |
| "input_ids": encoding["input_ids"].squeeze(), | |
| "attention_mask": encoding["attention_mask"].squeeze(), | |
| "labels": torch.tensor(self.labels[idx], dtype=torch.long), | |
| } | |
| # ───────────────────────────────────────────── | |
| # Transformer LID Model | |
| # ───────────────────────────────────────────── | |
| class TransformerLIDModel: | |
| """ | |
| Fine-tuned transformer for language identification. | |
| Defaults to MuRIL (designed for Indic languages). | |
| MuRIL advantages over XLM-R for this task: | |
| - Pretrained on 17 Indian languages + transliterated text | |
| - Better subword coverage for Indic scripts | |
| - Transliteration-aware pretraining | |
| """ | |
| SUPPORTED_MODELS = { | |
| "muril": "google/muril-base-cased", | |
| "xlmr": "xlm-roberta-base", | |
| "xlmr-large": "xlm-roberta-large", | |
| "indicbert": "ai4bharat/indic-bert", | |
| } | |
| def __init__( | |
| self, | |
| base_model: str = "muril", | |
| max_length: int = 128, | |
| batch_size: int = 32, | |
| learning_rate: float = 2e-5, | |
| num_epochs: int = 3, | |
| warmup_ratio: float = 0.1, | |
| weight_decay: float = 0.01, | |
| output_dir: str = "./artifacts/models/transformer", | |
| fp16: bool = True, | |
| ): | |
| assert TRANSFORMERS_AVAILABLE, "transformers + torch required" | |
| self.model_name = self.SUPPORTED_MODELS.get(base_model, base_model) | |
| self.max_length = max_length | |
| self.batch_size = batch_size | |
| self.learning_rate = learning_rate | |
| self.num_epochs = num_epochs | |
| self.warmup_ratio = warmup_ratio | |
| self.weight_decay = weight_decay | |
| self.output_dir = Path(output_dir) | |
| self.output_dir.mkdir(parents=True, exist_ok=True) | |
| self.fp16 = fp16 and torch.cuda.is_available() | |
| self.tokenizer = None | |
| self.model = None | |
| self.trainer = None | |
| self.is_fitted = False | |
| # Label mapping | |
| self.id2label = {i: l for i, l in LABEL_TO_CODE.items()} | |
| self.label2id = {l: i for i, l in self.id2label.items()} | |
| def _load_base_model(self): | |
| logger.info(f"Loading tokenizer and model: {self.model_name}") | |
| self.tokenizer = AutoTokenizer.from_pretrained( | |
| self.model_name, | |
| use_fast=True, | |
| ) | |
| self.model = AutoModelForSequenceClassification.from_pretrained( | |
| self.model_name, | |
| num_labels=NUM_CLASSES, | |
| id2label=self.id2label, | |
| label2id=self.label2id, | |
| ignore_mismatched_sizes=True, | |
| ) | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| logger.info(f"Using device: {device}") | |
| self.model.to(device) | |
| def fit( | |
| self, | |
| train_df: pd.DataFrame, | |
| val_df: Optional[pd.DataFrame] = None, | |
| ) -> "TransformerLIDModel": | |
| self._load_base_model() | |
| train_dataset = IndicLIDDataset(train_df, self.tokenizer, self.max_length) | |
| val_dataset = IndicLIDDataset(val_df, self.tokenizer, self.max_length) if val_df is not None else None | |
| training_args = TrainingArguments( | |
| output_dir=str(self.output_dir), | |
| num_train_epochs=self.num_epochs, | |
| per_device_train_batch_size=self.batch_size, | |
| per_device_eval_batch_size=self.batch_size * 2, | |
| learning_rate=self.learning_rate, | |
| warmup_ratio=self.warmup_ratio, | |
| weight_decay=self.weight_decay, | |
| evaluation_strategy="epoch" if val_dataset else "no", | |
| save_strategy="epoch", | |
| load_best_model_at_end=val_dataset is not None, | |
| metric_for_best_model="eval_loss", | |
| fp16=self.fp16, | |
| logging_steps=100, | |
| logging_dir=str(self.output_dir / "logs"), | |
| dataloader_num_workers=4, | |
| seed=42, | |
| report_to="none", # Disable wandb/mlflow unless configured | |
| ) | |
| self.trainer = Trainer( | |
| model=self.model, | |
| args=training_args, | |
| train_dataset=train_dataset, | |
| eval_dataset=val_dataset, | |
| tokenizer=self.tokenizer, | |
| callbacks=[EarlyStoppingCallback(early_stopping_patience=2)] if val_dataset else [], | |
| compute_metrics=self._compute_metrics, | |
| ) | |
| logger.info("Starting transformer training ...") | |
| self.trainer.train() | |
| self.is_fitted = True | |
| logger.success("Training complete!") | |
| return self | |
| def _compute_metrics(self, eval_pred): | |
| from sklearn.metrics import f1_score, accuracy_score | |
| logits, labels = eval_pred | |
| predictions = np.argmax(logits, axis=-1) | |
| acc = accuracy_score(labels, predictions) | |
| f1 = f1_score(labels, predictions, average="macro", zero_division=0) | |
| return {"accuracy": acc, "macro_f1": f1} | |
| def predict(self, texts: List[str]) -> List[str]: | |
| assert self.is_fitted | |
| probs, _ = self.predict_proba(texts) | |
| label_ids = np.argmax(probs, axis=1) | |
| return [self.id2label[i] for i in label_ids] | |
| def predict_proba(self, texts: List[str]) -> Tuple[np.ndarray, List[str]]: | |
| assert self.is_fitted | |
| import torch | |
| from torch.nn.functional import softmax | |
| device = next(self.model.parameters()).device | |
| all_probs = [] | |
| for i in range(0, len(texts), self.batch_size): | |
| batch_texts = texts[i:i + self.batch_size] | |
| encodings = self.tokenizer( | |
| batch_texts, | |
| max_length=self.max_length, | |
| truncation=True, | |
| padding=True, | |
| return_tensors="pt", | |
| ).to(device) | |
| with torch.no_grad(): | |
| outputs = self.model(**encodings) | |
| probs = softmax(outputs.logits, dim=-1).cpu().numpy() | |
| all_probs.append(probs) | |
| prob_matrix = np.vstack(all_probs) | |
| class_names = [self.id2label[i] for i in range(NUM_CLASSES)] | |
| return prob_matrix, class_names | |
| def predict_with_confidence(self, texts: List[str]) -> List[Dict]: | |
| probs, class_names = self.predict_proba(texts) | |
| results = [] | |
| for i, prob_row in enumerate(probs): | |
| top_k_idx = np.argsort(prob_row)[::-1][:5] | |
| top_k = [ | |
| { | |
| "lang": class_names[j], | |
| "name": LANGUAGE_REGISTRY.get(class_names[j], None) and LANGUAGE_REGISTRY[class_names[j]].name, | |
| "probability": float(prob_row[j]), | |
| } | |
| for j in top_k_idx | |
| ] | |
| best = top_k[0] | |
| info = LANGUAGE_REGISTRY.get(best["lang"]) | |
| results.append({ | |
| "text_preview": texts[i][:80], | |
| "predicted_lang": best["lang"], | |
| "predicted_lang_name": info.name if info else best["lang"], | |
| "script": info.script if info else "?", | |
| "direction": info.direction if info else "LTR", | |
| "confidence": best["probability"], | |
| "is_confident": best["probability"] >= 0.75, | |
| "top_5": top_k, | |
| }) | |
| return results | |
| def evaluate(self, test_df: pd.DataFrame) -> Dict: | |
| from sklearn.metrics import classification_report, f1_score, accuracy_score | |
| texts = test_df["text"].tolist() | |
| true_labels = test_df["lang"].tolist() | |
| pred_labels = self.predict(texts) | |
| probs, _ = self.predict_proba(texts) | |
| accuracy = accuracy_score(true_labels, pred_labels) | |
| macro_f1 = f1_score(true_labels, pred_labels, average="macro", zero_division=0) | |
| report = classification_report(true_labels, pred_labels, output_dict=True, zero_division=0) | |
| logger.info(f"\n{classification_report(true_labels, pred_labels, zero_division=0)}") | |
| logger.success(f"Accuracy: {accuracy:.4f} | Macro-F1: {macro_f1:.4f}") | |
| return { | |
| "accuracy": accuracy, | |
| "macro_f1": macro_f1, | |
| "per_language": report, | |
| "predictions": pred_labels, | |
| "true_labels": true_labels, | |
| "confidences": probs.max(axis=1).tolist(), | |
| } | |
| def save(self, path: str): | |
| self.model.save_pretrained(path) | |
| self.tokenizer.save_pretrained(path) | |
| logger.success(f"Transformer model saved: {path}") | |
| def load(cls, path: str) -> "TransformerLIDModel": | |
| assert TRANSFORMERS_AVAILABLE | |
| obj = cls.__new__(cls) | |
| obj.output_dir = Path(path) | |
| obj.tokenizer = AutoTokenizer.from_pretrained(path) | |
| obj.model = AutoModelForSequenceClassification.from_pretrained(path) | |
| obj.is_fitted = True | |
| obj.id2label = {i: l for i, l in LABEL_TO_CODE.items()} | |
| obj.batch_size = 32 | |
| obj.max_length = 128 | |
| obj.fp16 = False | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| obj.model.to(device) | |
| logger.success(f"Transformer model loaded: {path}") | |
| return obj | |