import base64 import os import re import shutil import subprocess import tempfile import zipfile from pathlib import Path from typing import Optional import numpy as np import soundfile as sf from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from huggingface_hub import hf_hub_download from pydantic import BaseModel import torch import torch.nn as nn import whisper from transformers import AutoModelForTokenClassification, AutoTokenizer, Wav2Vec2Model, Wav2Vec2Processor app = FastAPI(title="URTOX Toxic Span Detection API") MODEL_REPO_ID = "finalyear226/urdu-toxic-span-detector" MODEL_ZIP_NAME = "urtox_deploy_artifacts.zip" WHISPER_MODEL_SIZE = os.getenv("WHISPER_MODEL_SIZE", "small") ARTIFACTS_DIR = Path("artifacts") TEXT_MODEL_DIR = ARTIFACTS_DIR / "Urtox_attempt1" AUDIO_MODEL_PATH = ARTIFACTS_DIR / "audio_toxic_classifier.pt" LABELS_PATH = ARTIFACTS_DIR / "label_classes.npy" DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") TEXT_TOKENIZER = None TEXT_MODEL = None AUDIO_PROCESSOR = None AUDIO_WAV2VEC_MODEL = None AUDIO_CLASSIFIER = None AUDIO_LABELS = None WHISPER_MODEL = None MAX_AUDIO_LENGTH = 16000 * 10 URDU_PUNCTUATION = "،۔؟!؛:,.!?\"'()[]{}<>«»“”‘’" TOXIC_LEXICON = { "بہنچود", "بhenchod", "bhenchod", "بنچود", "مادرچود", "ماںچود", "چود", "چوتیا", "چوتیے", "چوتیئے", "حرامی", "حرامزادہ", "حرامزادی", "کنجر", "کنجری", "کمینہ", "کمینے", "بیوقوف", "احمق", "گھٹیا", "ذلیل", "خبیث", "بدتمیز", "بدتمیزی", "کتا", "کتے", "گدا", } app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) class DetectRequest(BaseModel): mode: str text: Optional[str] = None audio: Optional[str] = None class AudioToxicClassifier(nn.Module): def __init__(self, input_dim=768, hidden_dim=256, num_classes=2): super().__init__() self.classifier = nn.Sequential( nn.Linear(input_dim, hidden_dim), nn.ReLU(), nn.Dropout(0.3), nn.Linear(hidden_dim, 64), nn.ReLU(), nn.Dropout(0.2), nn.Linear(64, num_classes), ) def forward(self, x): return self.classifier(x) def artifacts_ready() -> bool: return ( TEXT_MODEL_DIR.exists() and (TEXT_MODEL_DIR / "model.safetensors").exists() and AUDIO_MODEL_PATH.exists() and LABELS_PATH.exists() ) def ensure_artifacts() -> None: if artifacts_ready(): return ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True) zip_path = hf_hub_download( repo_id=MODEL_REPO_ID, filename=MODEL_ZIP_NAME, repo_type="model", ) extract_dir = ARTIFACTS_DIR / "_downloaded" if extract_dir.exists(): shutil.rmtree(extract_dir) with zipfile.ZipFile(zip_path) as archive: archive.extractall(extract_dir) source_dir = extract_dir / "content" / "drive" / "MyDrive" if not source_dir.exists(): source_dir = extract_dir for name in ["Urtox_attempt1", "audio_toxic_classifier.pt", "label_classes.npy"]: source = source_dir / name destination = ARTIFACTS_DIR / name if destination.exists(): if destination.is_dir(): shutil.rmtree(destination) else: destination.unlink() if source.is_dir(): shutil.copytree(source, destination) else: shutil.copy2(source, destination) shutil.rmtree(extract_dir, ignore_errors=True) def load_text_model(): global TEXT_TOKENIZER, TEXT_MODEL if TEXT_TOKENIZER is not None and TEXT_MODEL is not None: return TEXT_TOKENIZER, TEXT_MODEL ensure_artifacts() TEXT_TOKENIZER = AutoTokenizer.from_pretrained(TEXT_MODEL_DIR) TEXT_MODEL = AutoModelForTokenClassification.from_pretrained(TEXT_MODEL_DIR) TEXT_MODEL.to(DEVICE) TEXT_MODEL.eval() return TEXT_TOKENIZER, TEXT_MODEL def load_audio_model(): global AUDIO_PROCESSOR, AUDIO_WAV2VEC_MODEL, AUDIO_CLASSIFIER, AUDIO_LABELS if ( AUDIO_PROCESSOR is not None and AUDIO_WAV2VEC_MODEL is not None and AUDIO_CLASSIFIER is not None and AUDIO_LABELS is not None ): return AUDIO_PROCESSOR, AUDIO_WAV2VEC_MODEL, AUDIO_CLASSIFIER, AUDIO_LABELS ensure_artifacts() AUDIO_PROCESSOR = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base") AUDIO_WAV2VEC_MODEL = Wav2Vec2Model.from_pretrained("facebook/wav2vec2-base") AUDIO_WAV2VEC_MODEL.to(DEVICE) AUDIO_WAV2VEC_MODEL.eval() AUDIO_LABELS = np.load(LABELS_PATH, allow_pickle=True).tolist() AUDIO_CLASSIFIER = AudioToxicClassifier(num_classes=len(AUDIO_LABELS)) AUDIO_CLASSIFIER.load_state_dict(torch.load(AUDIO_MODEL_PATH, map_location=DEVICE)) AUDIO_CLASSIFIER.to(DEVICE) AUDIO_CLASSIFIER.eval() return AUDIO_PROCESSOR, AUDIO_WAV2VEC_MODEL, AUDIO_CLASSIFIER, AUDIO_LABELS def load_whisper_model(): global WHISPER_MODEL if WHISPER_MODEL is not None: return WHISPER_MODEL WHISPER_MODEL = whisper.load_model(WHISPER_MODEL_SIZE, device=str(DEVICE)) return WHISPER_MODEL def normalize_word(word: str) -> str: normalized = word.strip().strip(URDU_PUNCTUATION).lower() normalized = re.sub(r"[\u064b-\u065f\u0670]", "", normalized) return normalized.replace(" ", "") def lexicon_match(word: str) -> bool: normalized = normalize_word(word) if not normalized: return False return normalized in TOXIC_LEXICON or any(term in normalized for term in TOXIC_LEXICON if len(term) >= 4) @app.on_event("startup") def startup_event(): ensure_artifacts() def predict_text(text: str): tokenizer, model = load_text_model() tokens = [token for token in text.split() if token] if not tokens: tokens = [" "] encoding = tokenizer( tokens, is_split_into_words=True, return_tensors="pt", truncation=True, max_length=128, padding="max_length", ) word_ids = encoding.word_ids(batch_index=0) model_inputs = {key: value.to(DEVICE) for key, value in encoding.items()} with torch.no_grad(): outputs = model(**model_inputs) probabilities = torch.softmax(outputs.logits, dim=-1)[0].cpu() predictions = torch.argmax(probabilities, dim=-1).tolist() id2label = model.config.id2label previous_word_id = None word_results = [] for token_index, word_id in enumerate(word_ids): if word_id is None or word_id == previous_word_id: continue model_label = id2label[int(predictions[token_index])] model_confidence = float(probabilities[token_index][predictions[token_index]]) fallback_toxic = lexicon_match(tokens[word_id]) label = model_label confidence = model_confidence if fallback_toxic and model_label == "O": label = "B-Toxic" confidence = max(model_confidence, 0.97) is_toxic = label in {"B-Toxic", "I-Toxic"} word_results.append( { "text": tokens[word_id], "toxic": is_toxic, "bioTag": label, "confidence": round(confidence, 4), "modelBioTag": model_label, "modelConfidence": round(model_confidence, 4), "source": "lexicon+model" if fallback_toxic and model_label == "O" else "model", } ) previous_word_id = word_id toxic_words = [word for word in word_results if word["toxic"]] toxic_confidences = [word["confidence"] for word in toxic_words] confidence = max(toxic_confidences) if toxic_confidences else 1.0 - max( (word["confidence"] for word in word_results), default=0.0, ) return { "isToxic": bool(toxic_words), "confidence": round(float(confidence), 4), "subLabel": "toxic" if toxic_words else "non-toxic", "subLabelConfidence": round(float(confidence), 4), "toxicSpanCount": count_toxic_spans(word_results), "transcript": None, "words": word_results, "xai": { "modelExplanation": "XLM-RoBERTa BIO token classification with a conservative Urdu abuse-word fallback for obvious missed slurs.", "topToxicTokens": [ { "token": word["text"], "attribution": word["confidence"], "confidence": word["confidence"], } for word in sorted(toxic_words, key=lambda item: item["confidence"], reverse=True)[:5] ], "integratedGradients": None, }, } def count_toxic_spans(words: list[dict]) -> int: span_count = 0 previous_toxic = False for word in words: current_toxic = bool(word["toxic"]) if current_toxic and not previous_toxic: span_count += 1 previous_toxic = current_toxic return span_count def decode_audio_to_tempfile(audio_payload: str) -> str: suffix = ".webm" if audio_payload.startswith("data:"): mime_type = audio_payload.split(";", 1)[0].replace("data:", "") if "webm" in mime_type: suffix = ".webm" elif "wav" in mime_type: suffix = ".wav" elif "mpeg" in mime_type or "mp3" in mime_type: suffix = ".mp3" elif "ogg" in mime_type: suffix = ".ogg" if "," in audio_payload: audio_payload = audio_payload.split(",", 1)[1] audio_bytes = base64.b64decode(audio_payload) temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=suffix) temp_file.write(audio_bytes) temp_file.close() return temp_file.name def convert_audio_to_wav(input_path: str) -> str: output_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav") output_file.close() command = [ "ffmpeg", "-y", "-i", input_path, "-ac", "1", "-ar", "16000", "-t", "10", output_file.name, ] subprocess.run(command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) return output_file.name def transcribe_audio(temp_path: str) -> str: whisper_model = load_whisper_model() result = whisper_model.transcribe( temp_path, language="ur", task="transcribe", fp16=DEVICE.type == "cuda", ) return (result.get("text") or "").strip() def predict_audio(audio_payload: str) -> dict: processor, wav2vec_model, audio_classifier, labels = load_audio_model() temp_path = decode_audio_to_tempfile(audio_payload) wav_path = None try: wav_path = convert_audio_to_wav(temp_path) transcript = transcribe_audio(wav_path) span_result = predict_text(transcript) if transcript else { "isToxic": False, "confidence": 0.0, "subLabel": "non-toxic", "subLabelConfidence": 0.0, "toxicSpanCount": 0, "transcript": None, "words": [], "xai": { "modelExplanation": "Whisper did not return a transcript for this audio.", "topToxicTokens": [], "integratedGradients": None, }, } waveform, sample_rate = sf.read(wav_path, dtype="float32") if waveform.ndim > 1: waveform = waveform.mean(axis=1) if waveform.shape[0] > MAX_AUDIO_LENGTH: waveform = waveform[:MAX_AUDIO_LENGTH] inputs = processor( waveform, sampling_rate=16000, return_tensors="pt", padding=True, ) inputs = {key: value.to(DEVICE) for key, value in inputs.items()} with torch.no_grad(): wav2vec_outputs = wav2vec_model(**inputs) embedding = wav2vec_outputs.last_hidden_state.mean(dim=1) logits = audio_classifier(embedding) probabilities = torch.softmax(logits, dim=1)[0].cpu().numpy() prediction_index = int(np.argmax(probabilities)) predicted_label = str(labels[prediction_index]) toxic_index = labels.index("toxic") if "toxic" in labels else prediction_index toxic_probability = float(probabilities[toxic_index]) confidence = float(probabilities[prediction_index]) is_audio_toxic = predicted_label == "toxic" is_toxic = is_audio_toxic or bool(span_result["isToxic"]) combined_confidence = max(confidence if is_audio_toxic else 0.0, float(span_result["confidence"])) return { "isToxic": is_toxic, "confidence": round(combined_confidence, 4), "subLabel": "toxic" if is_toxic else "non-toxic", "subLabelConfidence": round(combined_confidence, 4), "toxicSpanCount": span_result["toxicSpanCount"], "transcript": transcript, "words": span_result["words"], "audio": { "label": predicted_label, "toxicProbability": round(toxic_probability, 4), "nonToxicProbability": round(float(probabilities[labels.index("non_toxic")]), 4) if "non_toxic" in labels else None, }, "xai": { "modelExplanation": "Audio inference uses Whisper transcription for toxic-span detection plus Wav2Vec2 audio-level toxicity classification.", "topToxicTokens": span_result["xai"]["topToxicTokens"], "integratedGradients": span_result["xai"]["integratedGradients"], }, } finally: Path(temp_path).unlink(missing_ok=True) if wav_path: Path(wav_path).unlink(missing_ok=True) def audio_fallback_prediction(message: str = "Audio inference could not run.") -> dict: return { "isToxic": False, "confidence": 0.0, "subLabel": "audio-not-enabled", "subLabelConfidence": 0.0, "toxicSpanCount": 0, "transcript": message, "words": [], "xai": { "modelExplanation": message, "topToxicTokens": [], "integratedGradients": None, }, } @app.get("/") def health(): return { "status": "ok", "service": "urtox-api", "artifactSource": MODEL_REPO_ID, "artifactsReady": artifacts_ready(), "textModelLoaded": TEXT_MODEL is not None, "audioModelLoaded": AUDIO_CLASSIFIER is not None, "asrLoaded": WHISPER_MODEL is not None, "asrModel": f"openai-whisper/{WHISPER_MODEL_SIZE}", "device": str(DEVICE), } @app.post("/detect") def detect(payload: DetectRequest): if payload.mode == "audio": if not payload.audio: return audio_fallback_prediction("No audio payload was provided.") try: return predict_audio(payload.audio) except Exception as exc: return audio_fallback_prediction(f"Audio inference failed: {exc}") text = payload.text or "yeh toxic span detection result hai" return predict_text(text)