import os import re import io import json import time import zipfile from pathlib import Path from typing import List, Dict, Any, Tuple, Optional import numpy as np import pandas as pd import gradio as gr import torch from transformers import ( AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, GenerationConfig, ) # ========================= # Global config # ========================= SPACE_CACHE = Path.home() / ".cache" / "huggingface" SPACE_CACHE.mkdir(parents=True, exist_ok=True) DEVICE = "cuda" if torch.cuda.is_available() else "cpu" # Force slow tokenizer path by default; avoids Rust tokenizer.json parsing issues os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") os.environ.setdefault("TOKENIZERS_PREFER_FAST", "false") GEN_CONFIG = GenerationConfig( temperature=0.0, top_p=1.0, do_sample=False, max_new_tokens=128, # raise if JSON truncates ) OFFICIAL_LABELS = [ "plan_contact", "schedule_meeting", "update_contact_info_non_postal", "update_contact_info_postal_address", "update_kyc_activity", "update_kyc_origin_of_assets", "update_kyc_purpose_of_businessrelation", "update_kyc_total_assets", ] OFFICIAL_LABELS_TEXT = "\n".join(OFFICIAL_LABELS) # ========================= # Editable defaults (shown in UI) # ========================= DEFAULT_SYSTEM_INSTRUCTIONS = ( "You extract ACTIONABLE TASKS from client–advisor transcripts. " "The transcript may be in German, French, Italian, or English. " "Prioritize RECALL: if a label plausibly applies, include it. " "Use ONLY the canonical labels provided. " "Return STRICT JSON only with keys 'labels' and 'tasks'. " "Each task must include 'label', a brief 'explanation', and a short 'evidence' quote from the transcript." ) DEFAULT_LABEL_GLOSSARY = { "plan_contact": "Commitment to contact later (advisor/client will reach out, follow-up promised).", "schedule_meeting": "Scheduling or confirming a meeting/call/appointment (time/date/slot/virtual).", "update_contact_info_non_postal": "Change or confirmation of phone/email (non-postal contact details).", "update_contact_info_postal_address": "Change or confirmation of postal/residential/mailing address.", "update_kyc_activity": "Change/confirmation of occupation, employment status, or economic activity.", "update_kyc_origin_of_assets": "Discussion/confirmation of source of funds / origin of assets.", "update_kyc_purpose_of_businessrelation": "Purpose of the banking relationship/account usage.", "update_kyc_total_assets": "Discussion/confirmation of total assets/net worth.", } # Tiny multilingual fallback rules (optional) to avoid empty outputs DEFAULT_FALLBACK_CUES = { "plan_contact": [ r"\b(get|got|will|we'?ll|i'?ll)\s+back to you\b", r"\bfollow\s*up\b", r"\breach out\b", r"\btouch base\b", r"\bcontact (you|me|us)\b", r"\bin verbindung setzen\b", r"\brückmeldung\b", r"\bich\s+melde\b|\bwir\s+melden\b", r"\bnachfassen\b", r"\bje vous recontacte\b|\bnous vous recontacterons\b", r"\bprendre contact\b|\breprendre contact\b", r"\bla ricontatter[oò]\b|\bci metteremo in contatto\b", r"\btenersi in contatto\b", ], "schedule_meeting": [ r"\b(let'?s\s+)?meet(ing|s)?\b", r"\bschedule( a)? (call|meeting|appointment)\b", r"\bbook( a)? (slot|time|meeting)\b", r"\b(next week|tomorrow|this (afternoon|morning|evening))\b", r"\bconfirm( the)? (time|meeting|appointment)\b", r"\btermin(e|s)?\b|\bvereinbaren\b|\bansetzen\b|\babstimmen\b|\bbesprechung(en)?\b|\bvirtuell(e|en)?\b", r"\bnächste(n|r)? woche\b|\b(dienstag|montag|mittwoch|donnerstag|freitag)\b|\bnachmittag|vormittag|morgen\b", r"\brendez[- ]?vous\b|\bréunion\b|\bfixer\b|\bplanifier\b|\bse rencontrer\b|\bse voir\b", r"\bla semaine prochaine\b|\bdemain\b|\bcet (après-midi|apres-midi|après midi|apres midi|matin|soir)\b", r"\bappuntamento\b|\briunione\b|\borganizzare\b|\bprogrammare\b|\bincontrarci\b|\bcalendario\b", r"\bla prossima settimana\b|\bdomani\b|\b(questo|questa)\s*(pomeriggio|mattina|sera)\b", ], "update_kyc_origin_of_assets": [ r"\bsource of funds\b|\borigin of assets\b|\bproof of (funds|assets)\b", r"\bvermögensursprung(e|s)?\b|\bherkunft der mittel\b|\bnachweis\b", r"\borigine des fonds\b|\borigine du patrimoine\b|\bjustificatif(s)?\b", r"\borigine dei fondi\b|\borigine del patrimonio\b|\bprova dei fondi\b|\bgiustificativo\b", ], "update_kyc_activity": [ r"\bemployment status\b|\boccupation\b|\bjob change\b|\bsalary history\b", r"\bbeschäftigungsstatus\b|\bberuf\b|\bjobwechsel\b|\bgehaltshistorie\b|\btätigkeit\b", r"\bstatut professionnel\b|\bprofession\b|\bchangement d'emploi\b|\bhistorique salarial\b|\bactivité\b", r"\bstato occupazionale\b|\bprofessione\b|\bcambio di lavoro\b|\bstoria salariale\b|\battivit[aà]\b", ], } # ========================= # Prompt template # ========================= USER_PROMPT_TEMPLATE = ( "Transcript (may be DE/FR/IT/EN):\n" "```\n{transcript}\n```\n\n" "Allowed Labels (canonical; use only these):\n" "{allowed_labels_list}\n\n" "Label Glossary (concise semantics):\n" "{glossary}\n\n" "Return STRICT JSON ONLY in this exact schema:\n" '{\n "labels": ["", "..."],\n' ' "tasks": [{"label": "", "explanation": "", "evidence": ""}]\n}\n' ) # ========================= # Utilities # ========================= def _now_ms() -> int: return int(time.time() * 1000) def normalize_labels(labels: List[str]) -> List[str]: return list(dict.fromkeys([l.strip() for l in labels if isinstance(l, str) and l.strip()])) def canonicalize_map(allowed: List[str]) -> Dict[str, str]: return {lab.lower(): lab for lab in allowed} def robust_json_extract(text: str) -> Dict[str, Any]: if not text: return {"labels": [], "tasks": []} start, end = text.find("{"), text.rfind("}") candidate = text[start:end+1] if (start != -1 and end != -1 and end > start) else text try: return json.loads(candidate) except Exception: candidate = re.sub(r",\s*}", "}", candidate) candidate = re.sub(r",\s*]", "]", candidate) try: return json.loads(candidate) except Exception: return {"labels": [], "tasks": []} def restrict_to_allowed(pred: Dict[str, Any], allowed: List[str]) -> Dict[str, Any]: out = {"labels": [], "tasks": []} allowed_map = canonicalize_map(allowed) filt_labels = [] for l in pred.get("labels", []) or []: k = str(l).strip().lower() if k in allowed_map: filt_labels.append(allowed_map[k]) filt_labels = normalize_labels(filt_labels) filt_tasks = [] for t in pred.get("tasks", []) or []: if not isinstance(t, dict): continue k = str(t.get("label", "")).strip().lower() if k in allowed_map: new_t = dict(t); new_t["label"] = allowed_map[k] new_t = { "label": new_t["label"], "explanation": str(new_t.get("explanation", ""))[:300], "evidence": str(new_t.get("evidence", ""))[:300], } filt_tasks.append(new_t) merged = normalize_labels(list(set(filt_labels) | {tt["label"] for tt in filt_tasks})) out["labels"] = merged out["tasks"] = filt_tasks return out # ========================= # Pre-processing # ========================= _DISCLAIMER_PATTERNS = [ r"(?is)^\s*(?:disclaimer|legal notice|confidentiality notice).+?(?:\n{2,}|$)", r"(?is)^\s*the information contained.+?(?:\n{2,}|$)", r"(?is)^\s*this message \(including any attachments\).+?(?:\n{2,}|$)", ] _FOOTER_PATTERNS = [ r"(?is)\n+kind regards[^\n]*\n.*$", r"(?is)\n+best regards[^\n]*\n.*$", r"(?is)\n+sent from my.*$", r"(?is)\n+ubs ag.*$", ] _TIMESTAMP_SPEAKER = [ r"\[\d{1,2}:\d{2}(:\d{2})?\]", r"^\s*(advisor|client|client advisor)\s*:\s*", r"^\s*(speaker\s*\d+)\s*:\s*", ] def clean_transcript(text: str) -> str: if not text: return text s = text # strip speaker/timestamps lines = [] for ln in s.splitlines(): ln2 = ln for pat in _TIMESTAMP_SPEAKER: ln2 = re.sub(pat, "", ln2, flags=re.IGNORECASE) lines.append(ln2) s = "\n".join(lines) # disclaimers (top) for pat in _DISCLAIMER_PATTERNS: s = re.sub(pat, "", s).strip() # footers for pat in _FOOTER_PATTERNS: s = re.sub(pat, "", s) # whitespace tidy s = re.sub(r"[ \t]+", " ", s) s = re.sub(r"\n{3,}", "\n\n", s).strip() return s def read_text_file_any(file_input) -> str: if not file_input: return "" if isinstance(file_input, (str, Path)): try: return Path(file_input).read_text(encoding="utf-8", errors="ignore") except Exception: return "" try: data = file_input.read() return data.decode("utf-8", errors="ignore") except Exception: return "" def read_json_file_any(file_input) -> Optional[dict]: if not file_input: return None if isinstance(file_input, (str, Path)): try: return json.loads(Path(file_input).read_text(encoding="utf-8", errors="ignore")) except Exception: return None try: return json.loads(file_input.read().decode("utf-8", errors="ignore")) except Exception: return None def truncate_tokens(tokenizer, text: str, max_tokens: int) -> str: toks = tokenizer(text, add_special_tokens=False)["input_ids"] if len(toks) <= max_tokens: return text return tokenizer.decode(toks[-max_tokens:], skip_special_tokens=True) # ========================= # Cache purge for fresh downloads # ========================= def _purge_repo_from_cache(repo_id: str): """Delete cached files of a specific repo to guarantee a fresh download.""" try: base = SPACE_CACHE safe = repo_id.replace("/", "--") for p in base.glob(f"models--{safe}*"): try: if p.is_file(): p.unlink() else: for sub in sorted(p.rglob("*"), reverse=True): try: if sub.is_file() or sub.is_symlink(): sub.unlink() else: sub.rmdir() except Exception: pass p.rmdir() except Exception: pass except Exception: pass # ========================= # HF model wrapper (robust: slow tokenizer first + load fallbacks) # ========================= class ModelWrapper: def __init__(self, repo_id: str, hf_token: Optional[str], load_in_4bit: bool, use_sdpa: bool, force_tok_redownload: bool): self.repo_id = repo_id self.hf_token = hf_token self.load_in_4bit = load_in_4bit self.use_sdpa = use_sdpa self.force_tok_redownload = force_tok_redownload self.tokenizer = None self.model = None self.load_path = "uninitialized" def _load_tokenizer(self): """ Prefer the slow (SentencePiece) tokenizer first to avoid Rust tokenizers JSON parsing. If user asked to force fresh download, purge local cache first. """ if self.force_tok_redownload: _purge_repo_from_cache(self.repo_id) common = dict( pretrained_model_name_or_path=self.repo_id, token=self.hf_token, cache_dir=str(SPACE_CACHE), trust_remote_code=True, local_files_only=False, force_download=True if self.force_tok_redownload else False, revision=None, ) # 1) SLOW PATH FIRST slow_err = None tok = None try: tok = AutoTokenizer.from_pretrained(use_fast=False, **common) except Exception as e: slow_err = e # 2) If slow somehow failed, try FAST as a last resort fast_err = None if tok is None: try: tok = AutoTokenizer.from_pretrained(use_fast=True, **common) except Exception as e: fast_err = e if tok is None: raise RuntimeError(f"Tokenizer failed (slow: {slow_err}) (fast: {fast_err})") if tok.pad_token is None and tok.eos_token: tok.pad_token = tok.eos_token # Tag which path we used if slow_err is None: self.load_path = "tok:SLOW" else: self.load_path = "tok:FAST" return tok def load(self): qcfg = None if self.load_in_4bit and DEVICE == "cuda": qcfg = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.float16, bnb_4bit_use_double_quant=True, ) tok = self._load_tokenizer() errors = [] for desc, kwargs in [ ("auto_device_no_lowcpu" + ("_sdpa" if (self.use_sdpa and DEVICE=="cuda") else ""), dict( torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32, device_map="auto" if DEVICE == "cuda" else None, low_cpu_mem_usage=False, quantization_config=qcfg, trust_remote_code=True, cache_dir=str(SPACE_CACHE), attn_implementation=("sdpa" if (self.use_sdpa and DEVICE == "cuda") else None), )), ("auto_device_no_sdpa", dict( torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32, device_map="auto" if DEVICE == "cuda" else None, low_cpu_mem_usage=False, quantization_config=qcfg, trust_remote_code=True, cache_dir=str(SPACE_CACHE), )), ("cpu_then_to_cuda" if DEVICE == "cuda" else "cpu_only", dict( torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32, device_map=None, low_cpu_mem_usage=False, quantization_config=None if DEVICE != "cuda" else qcfg, trust_remote_code=True, cache_dir=str(SPACE_CACHE), )), ]: try: mdl = AutoModelForCausalLM.from_pretrained(self.repo_id, token=self.hf_token, **kwargs) if desc.startswith("cpu_then_to_cuda") and DEVICE == "cuda": mdl = mdl.to(torch.device("cuda")) self.tokenizer = tok self.model = mdl self.load_path = f"{self.load_path} | {desc}" return except Exception as e: errors.append(f"{desc}: {e}") raise RuntimeError("All load attempts failed:\n" + "\n".join(errors)) @torch.inference_mode() def generate(self, system_prompt: str, user_prompt: str) -> str: if hasattr(self.tokenizer, "apply_chat_template"): messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_prompt}, ] input_ids = self.tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_tensors="pt" ) input_ids = input_ids.to(self.model.device) gen_kwargs = dict( input_ids=input_ids, generation_config=GEN_CONFIG, eos_token_id=self.tokenizer.eos_token_id, pad_token_id=self.tokenizer.pad_token_id, ) else: enc = self.tokenizer( f"[SYSTEM]\n{system_prompt}\n[/SYSTEM]\n[USER]\n{user_prompt}\n[/USER]\n", return_tensors="pt" ).to(self.model.device) gen_kwargs = dict( **enc, generation_config=GEN_CONFIG, eos_token_id=self.tokenizer.eos_token_id, pad_token_id=self.tokenizer.pad_token_id, ) with torch.cuda.amp.autocast(enabled=(DEVICE == "cuda")): out_ids = self.model.generate(**gen_kwargs) return self.tokenizer.decode(out_ids[0], skip_special_tokens=True) _MODEL_CACHE: Dict[str, ModelWrapper] = {} def get_model(repo_id: str, hf_token: Optional[str], load_in_4bit: bool, use_sdpa: bool, force_tok_redownload: bool) -> ModelWrapper: key = f"{repo_id}::{'4bit' if (load_in_4bit and DEVICE=='cuda') else 'full'}::{'sdpa' if use_sdpa else 'nosdpa'}::{'force' if force_tok_redownload else 'cache'}" if key not in _MODEL_CACHE: m = ModelWrapper(repo_id, hf_token, load_in_4bit, use_sdpa, force_tok_redownload) m.load() _MODEL_CACHE[key] = m return _MODEL_CACHE[key] # ========================= # Evaluation (official weighted score) # ========================= def evaluate_predictions(y_true: List[List[str]], y_pred: List[List[str]]) -> float: ALLOWED_LABELS = OFFICIAL_LABELS LABEL_TO_IDX = {label: idx for idx, label in enumerate(ALLOWED_LABELS)} def _process_sample_labels(sample_labels: List[str], sample_name: str) -> List[str]: if not isinstance(sample_labels, list): raise ValueError(f"{sample_name} must be a list of strings, got {type(sample_labels)}") seen, uniq = set(), [] for label in sample_labels: if not isinstance(label, str): raise ValueError(f"{sample_name} contains non-string: {label} (type: {type(label)})") if label in seen: raise ValueError(f"{sample_name} contains duplicate label: '{label}'") if label not in ALLOWED_LABELS: raise ValueError(f"{sample_name} contains invalid label: '{label}'. Allowed: {ALLOWED_LABELS}") seen.add(label); uniq.append(label) return uniq if len(y_true) != len(y_pred): raise ValueError(f"y_true and y_pred must have same length. Got {len(y_true)} vs {len(y_pred)}") n_samples = len(y_true) n_labels = len(OFFICIAL_LABELS) y_true_binary = np.zeros((n_samples, n_labels), dtype=int) y_pred_binary = np.zeros((n_samples, n_labels), dtype=int) for i, sample_labels in enumerate(y_true): for label in _process_sample_labels(sample_labels, f"y_true[{i}]"): y_true_binary[i, LABEL_TO_IDX[label]] = 1 for i, sample_labels in enumerate(y_pred): for label in _process_sample_labels(sample_labels, f"y_pred[{i}]"): y_pred_binary[i, LABEL_TO_IDX[label]] = 1 fn = np.sum((y_true_binary == 1) & (y_pred_binary == 0), axis=1) fp = np.sum((y_true_binary == 0) & (y_pred_binary == 1), axis=1) weighted = 2.0 * fn + 1.0 * fp max_err = 2.0 * np.sum(y_true_binary, axis=1) + 1.0 * (n_labels - np.sum(y_true_binary, axis=1)) per_sample = np.where(max_err > 0, 1.0 - (weighted / max_err), 1.0) return float(max(0.0, min(1.0, np.mean(per_sample)))) # ========================= # Multilingual regex fallback (optional) # ========================= def multilingual_fallback(text: str, allowed: List[str], cues: Dict[str, List[str]]) -> Dict[str, Any]: low = text.lower() labels, tasks = [], [] for lab in allowed: for pat in cues.get(lab, []): m = re.search(pat, low) if m: i = m.start() start = max(0, i - 60); end = min(len(text), i + len(m.group(0)) + 60) if lab not in labels: labels.append(lab) tasks.append({ "label": lab, "explanation": "Rule hit (multilingual fallback)", "evidence": text[start:end].strip() }) break return {"labels": normalize_labels(labels), "tasks": tasks} # ========================= # Inference helpers # ========================= def build_glossary_str(glossary: Dict[str, str], allowed: List[str]) -> str: return "\n".join([f"- {lab}: {glossary.get(lab, '')}" for lab in allowed]) def warmup_model(model_repo: str, use_4bit: bool, use_sdpa: bool, hf_token: str, force_tok_redownload: bool) -> str: t0 = _now_ms() try: model = get_model(model_repo, (hf_token or "").strip() or None, use_4bit, use_sdpa, force_tok_redownload) _ = model.generate("Return JSON only.", '{"labels": [], "tasks": []}') return f"Warm-up complete in {_now_ms() - t0} ms. Load path: {model.load_path}" except Exception as e: return f"Warm-up failed: {e}" def run_single( transcript_text: str, transcript_file, gt_json_text: str, gt_json_file, use_cleaning: bool, use_fallback: bool, allowed_labels_text: str, sys_instructions_text: str, glossary_json_text: str, fallback_json_text: str, model_repo: str, use_4bit: bool, use_sdpa: bool, max_input_tokens: int, hf_token: str, force_tok_redownload: bool, ) -> Tuple[str, str, str, str, str, str, str, str, str]: t0 = _now_ms() raw_text = "" if transcript_file: raw_text = read_text_file_any(transcript_file) raw_text = (raw_text or transcript_text or "").strip() if not raw_text: return "", "", "No transcript provided.", "", "", "", "", "", "" text = clean_transcript(raw_text) if use_cleaning else raw_text user_allowed = [ln.strip() for ln in (allowed_labels_text or "").splitlines() if ln.strip()] allowed = normalize_labels(user_allowed or OFFICIAL_LABELS) try: sys_instructions = (sys_instructions_text or DEFAULT_SYSTEM_INSTRUCTIONS).strip() or DEFAULT_SYSTEM_INSTRUCTIONS except Exception: sys_instructions = DEFAULT_SYSTEM_INSTRUCTIONS try: label_glossary = json.loads(glossary_json_text) if glossary_json_text else DEFAULT_LABEL_GLOSSARY except Exception: label_glossary = DEFAULT_LABEL_GLOSSARY try: fallback_cues = json.loads(fallback_json_text) if fallback_json_text else DEFAULT_FALLBACK_CUES except Exception: fallback_cues = DEFAULT_FALLBACK_CUES try: model = get_model(model_repo, (hf_token or "").strip() or None, use_4bit, use_sdpa, force_tok_redownload) except Exception as e: return "", "", f"Model load failed: {e}", "", "", "", "", "", "" trunc = truncate_tokens(model.tokenizer, text, max_input_tokens) glossary_str = build_glossary_str(label_glossary, allowed) allowed_list_str = "\n".join(f"- {l}" for l in allowed) user_prompt = USER_PROMPT_TEMPLATE.format( transcript=trunc, allowed_labels_list=allowed_list_str, glossary=glossary_str, ) transcript_tokens = len(model.tokenizer(trunc, add_special_tokens=False)["input_ids"]) prompt_tokens = len(model.tokenizer(user_prompt, add_special_tokens=False)["input_ids"]) token_info_text = f"Transcript tokens: {transcript_tokens} | Prompt tokens: {prompt_tokens} | Load path: {model.load_path}" prompt_preview_text = "```\n" + user_prompt[:4000] + ("\n... (truncated)" if len(user_prompt) > 4000 else "") + "\n```" t1 = _now_ms() try: out = model.generate(sys_instructions, user_prompt) except Exception as e: return "", "", f"Generation error: {e}", "", "", "", prompt_preview_text, token_info_text, "" t2 = _now_ms() parsed = robust_json_extract(out) filtered = restrict_to_allowed(parsed, allowed) if use_fallback: fb = multilingual_fallback(trunc, allowed, fallback_cues) if fb["labels"]: merged_labels = sorted(list(set(filtered.get("labels", [])) | set(fb["labels"]))) existing = {tt.get("label") for tt in filtered.get("tasks", [])} merged_tasks = filtered.get("tasks", []) + [t for t in fb["tasks"] if t["label"] not in existing] filtered = {"labels": merged_labels, "tasks": merged_tasks} diag = "\n".join([ f"Device: {DEVICE} (4-bit: {'Yes' if (use_4bit and DEVICE=='cuda') else 'No'})", f"Model: {model_repo}", f"Input cleaned: {'Yes' if use_cleaning else 'No'}", f"Fallback rules: {'Yes' if use_fallback else 'No'}", f"SDPA attention: {'Yes' if use_sdpa else 'No'}", f"Tokens (input limit): ≤ {max_input_tokens}", f"Latency: prep {t1-t0} ms, gen {t2-t1} ms, total {t2-t0} ms", f"Allowed labels: {', '.join(allowed)}", ]) labs = filtered.get("labels", []) tasks = filtered.get("tasks", []) summary = "Detected labels:\n" + ("\n".join(f"- {l}" for l in labs) if labs else "(none)") if tasks: summary += "\n\nTasks:\n" + "\n".join( f"• [{t['label']}] {t.get('explanation','')} | ev: {t.get('evidence','')[:140]}{'…' if len(t.get('evidence',''))>140 else ''}" for t in tasks ) else: summary += "\n\nTasks: (none)" json_out = json.dumps(filtered, indent=2, ensure_ascii=False) metrics = "" if gt_json_file or (gt_json_text and gt_json_text.strip()): truth_obj = None if gt_json_file: truth_obj = read_json_file_any(gt_json_file) if (not truth_obj) and gt_json_text: try: truth_obj = json.loads(gt_json_text) except Exception: pass if isinstance(truth_obj, dict) and isinstance(truth_obj.get("labels"), list): true_labels = [x for x in truth_obj["labels"] if x in OFFICIAL_LABELS] pred_labels = labs try: score = evaluate_predictions([true_labels], [pred_labels]) tp = len(set(true_labels) & set(pred_labels)) fp = len(set(pred_labels) - set(true_labels)) fn = len(set(true_labels) - set(pred_labels)) recall = tp / (tp + fn) if (tp + fn) else 1.0 precision = tp / (tp + fp) if (tp + fp) else 1.0 f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 1.0 metrics = ( f"Weighted score: {score:.3f}\n" f"Recall: {recall:.3f} | Precision: {precision:.3f} | F1: {f1:.3f}\n" f"TP={tp} FP={fp} FN={fn}\n" f"Truth: {', '.join(true_labels)}" ) except Exception as e: metrics = f"Scoring error: {e}" else: metrics = "Ground truth JSON missing or invalid; expected {'labels': [...]}." context_preview = "### Label Glossary (used)\n" + "\n".join(f"- {k}: {v}" for k, v in DEFAULT_LABEL_GLOSSARY.items() if k in allowed) instructions_preview = "```\n" + (sys_instructions_text or DEFAULT_SYSTEM_INSTRUCTIONS) + "\n```" return summary, json_out, diag, out.strip(), context_preview, instructions_preview, metrics, prompt_preview_text, token_info_text # ========================= # Batch mode # ========================= def read_zip_from_path(path: str, exdir: Path) -> List[Path]: exdir.mkdir(parents=True, exist_ok=True) with open(path, "rb") as f: data = f.read() with zipfile.ZipFile(io.BytesIO(data)) as zf: zf.extractall(exdir) return [p for p in exdir.rglob("*") if p.is_file()] def run_batch( zip_path, use_cleaning: bool, use_fallback: bool, sys_instructions_text: str, glossary_json_text: str, fallback_json_text: str, model_repo: str, use_4bit: bool, use_sdpa: bool, max_input_tokens: int, hf_token: str, force_tok_redownload: bool, limit_files: int, ) -> Tuple[str, str, pd.DataFrame, str]: if not zip_path: return ("No ZIP provided.", "", pd.DataFrame(), "") try: sys_instructions = (sys_instructions_text or DEFAULT_SYSTEM_INSTRUCTIONS).strip() or DEFAULT_SYSTEM_INSTRUCTIONS except Exception: sys_instructions = DEFAULT_SYSTEM_INSTRUCTIONS try: label_glossary = json.loads(glossary_json_text) if glossary_json_text else DEFAULT_LABEL_GLOSSARY except Exception: label_glossary = DEFAULT_LABEL_GLOSSARY try: fallback_cues = json.loads(fallback_json_text) if fallback_json_text else DEFAULT_FALLBACK_CUES except Exception: fallback_cues = DEFAULT_FALLBACK_CUES work = Path("/tmp/batch") if work.exists(): for p in sorted(work.rglob("*"), reverse=True): try: p.unlink() except Exception: pass try: work.rmdir() except Exception: pass work.mkdir(parents=True, exist_ok=True) files = read_zip_from_path(zip_path, work) txts: Dict[str, Path] = {} gts: Dict[str, Path] = {} for p in files: if p.suffix.lower() == ".txt": txts[p.stem] = p elif p.suffix.lower() == ".json": gts[p.stem] = p stems = sorted(txts.keys()) if limit_files > 0: stems = stems[:limit_files] if not stems: return ("No .txt transcripts found in ZIP.", "", pd.DataFrame(), "") try: model = get_model(model_repo, (hf_token or "").strip() or None, use_4bit, use_sdpa, force_tok_redownload) except Exception as e: return (f"Model load failed: {e}", "", pd.DataFrame(), "") allowed = OFFICIAL_LABELS[:] glossary_str = build_glossary_str(label_glossary, allowed) allowed_list_str = "\n".join(f"- {l}" for l in allowed) y_true, y_pred = [], [] rows = [] t_start = _now_ms() for stem in stems: raw = txts[stem].read_text(encoding="utf-8", errors="ignore") text = clean_transcript(raw) if use_cleaning else raw trunc = truncate_tokens(model.tokenizer, text, max_input_tokens) user_prompt = USER_PROMPT_TEMPLATE.format( transcript=trunc, allowed_labels_list=allowed_list_str, glossary=glossary_str, ) t0 = _now_ms() out = model.generate(sys_instructions, user_prompt) t1 = _now_ms() parsed = robust_json_extract(out) filtered = restrict_to_allowed(parsed, allowed) if use_fallback: fb = multilingual_fallback(trunc, allowed, fallback_cues) if fb["labels"]: merged_labels = sorted(list(set(filtered.get("labels", [])) | set(fb["labels"]))) existing = {tt.get("label") for tt in filtered.get("tasks", [])} merged_tasks = filtered.get("tasks", []) + [t for t in fb["tasks"] if t["label"] not in existing] filtered = {"labels": merged_labels, "tasks": merged_tasks} pred_labels = filtered.get("labels", []) y_pred.append(pred_labels) gt_labels = [] if stem in gts: try: gt_obj = json.loads(gts[stem].read_text(encoding="utf-8", errors="ignore")) if isinstance(gt_obj, dict) and isinstance(gt_obj.get("labels"), list): gt_labels = [x for x in gt_obj["labels"] if x in OFFICIAL_LABELS] except Exception: pass y_true.append(gt_labels) gt_set, pr_set = set(gt_labels), set(pred_labels) tp = sorted(gt_set & pr_set) fp = sorted(pr_set - gt_set) fn = sorted(gt_set - pr_set) rows.append({ "file": stem, "true_labels": ", ".join(gt_labels), "pred_labels": ", ".join(pred_labels), "TP": len(tp), "FP": len(fp), "FN": len(fn), "gen_ms": t1 - t0 }) have_truth = any(len(v) > 0 for v in y_true) score = evaluate_predictions(y_true, y_pred) if have_truth else None df = pd.DataFrame(rows).sort_values(["FN", "FP", "file"]) diag = [ f"Processed files: {len(stems)}", f"Device: {DEVICE} (4-bit: {'Yes' if (use_4bit and DEVICE=='cuda') else 'No'})", f"Model: {model_repo}", f"Fallback rules: {'Yes' if use_fallback else 'No'}", f"SDPA attention: {'Yes' if use_sdpa else 'No'}", f"Tokens (input limit): ≤ {max_input_tokens}", f"Batch time: {_now_ms()-t_start} ms", ] if have_truth and score is not None: total_tp = int(df["TP"].sum()) total_fp = int(df["FP"].sum()) total_fn = int(df["FN"].sum()) recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) else 1.0 precision = total_tp / (total_tp + total_fp) if (total_tp + total_fp) else 1.0 f1 = (2 * precision * recall / (precision + recall)) if (precision + recall) else 1.0 diag += [ f"Official weighted score (0–1): {score:.3f}", f"Recall: {recall:.3f} | Precision: {precision:.3f} | F1: {f1:.3f}", f"Total TP={total_tp} FP={total_fp} FN={total_fn}", ] diag_str = "\n".join(diag) out_csv = Path("/tmp/batch_results.csv") df.to_csv(out_csv, index=False, encoding="utf-8") return ("Batch done.", diag_str, df, str(out_csv)) # ========================= # UI # ========================= MODEL_CHOICES = [ "swiss-ai/Apertus-8B-Instruct-2509", "meta-llama/Meta-Llama-3-8B-Instruct", "mistralai/Mistral-7B-Instruct-v0.3", ] # White, modern UI (no purple) custom_css = """ :root { --radius: 14px; } .gradio-container { font-family: Inter, ui-sans-serif, system-ui; background: #ffffff; color: #111827; } .card { border: 1px solid #e5e7eb; border-radius: var(--radius); padding: 14px 16px; background: #ffffff; box-shadow: 0 1px 2px rgba(0,0,0,.03); } .header { font-weight: 700; font-size: 22px; margin-bottom: 4px; color: #0f172a; } .subtle { color: #475569; font-size: 14px; margin-bottom: 12px; } hr.sep { border: none; border-top: 1px solid #e5e7eb; margin: 10px 0 16px; } .gr-button { border-radius: 12px !important; } a, .prose a { color: #0ea5e9; } """ with gr.Blocks(theme=gr.themes.Soft(), css=custom_css, fill_height=True) as demo: gr.Markdown("
Talk2Task — Multilingual Task Extraction (UBS Challenge)
") gr.Markdown("
Single-pass multilingual extraction (DE/FR/IT/EN). Optional rules fallback for recall. Batch evaluation included.
") with gr.Tab("Single transcript"): with gr.Row(): with gr.Column(scale=3): gr.Markdown("
Transcript
") file = gr.File( label="Drag & drop transcript (.txt / .md / .json)", file_types=[".txt", ".md", ".json"], type="filepath", ) text = gr.Textbox(label="Or paste transcript", lines=10, placeholder="Paste transcript in DE/FR/IT/EN…") gr.Markdown("
") gr.Markdown("
Ground truth JSON (optional)
") gt_file = gr.File( label="Upload ground truth JSON (expects {'labels': [...]})", file_types=[".json"], type="filepath", ) gt_text = gr.Textbox(label="Or paste ground truth JSON", lines=6, placeholder='{\"labels\": [\"schedule_meeting\"]}') gr.Markdown("
") # close card gr.Markdown("
Processing options
") use_cleaning = gr.Checkbox(label="Apply default cleaning (remove disclaimers, timestamps, speakers, footers)", value=True) use_fallback = gr.Checkbox(label="Enable multilingual fallback rule layer", value=True) gr.Markdown("
") gr.Markdown("
Allowed labels
") labels_text = gr.Textbox(label="Allowed Labels (one per line)", value=OFFICIAL_LABELS_TEXT, lines=8) reset_btn = gr.Button("Reset to official labels") gr.Markdown("
") gr.Markdown("
Editable instructions & context
") sys_instr_tb = gr.Textbox(label="System Instructions (editable)", value=DEFAULT_SYSTEM_INSTRUCTIONS, lines=5) glossary_tb = gr.Code(label="Label Glossary (JSON; editable)", value=json.dumps(DEFAULT_LABEL_GLOSSARY, indent=2), language="json") fallback_tb = gr.Code(label="Fallback Cues (Multilingual, JSON; editable)", value=json.dumps(DEFAULT_FALLBACK_CUES, indent=2), language="json") gr.Markdown("
") with gr.Column(scale=2): gr.Markdown("
Model & run
") repo = gr.Dropdown(label="Model", choices=MODEL_CHOICES, value=MODEL_CHOICES[0]) use_4bit = gr.Checkbox(label="Use 4-bit (GPU only)", value=True) use_sdpa = gr.Checkbox(label="Use SDPA attention (faster on many GPUs)", value=True) force_tok_redownload = gr.Checkbox(label="Force fresh tokenizer download", value=False) max_tokens = gr.Slider(label="Max input tokens", minimum=1024, maximum=8192, step=512, value=2048) hf_token = gr.Textbox(label="HF_TOKEN (only for gated models)", type="password", value=os.environ.get("HF_TOKEN","")) warm_btn = gr.Button("Warm up model (load & compile kernels)") run_btn = gr.Button("Run Extraction", variant="primary") gr.Markdown("
") gr.Markdown("
Outputs
") summary = gr.Textbox(label="Summary", lines=12) json_out = gr.Code(label="Strict JSON Output", language="json") diag = gr.Textbox(label="Diagnostics", lines=10) raw = gr.Textbox(label="Raw Model Output", lines=8) metrics_tb = gr.Textbox(label="Metrics vs Ground Truth (optional)", lines=6) prompt_preview = gr.Code(label="Prompt preview (user prompt sent)", language="markdown") token_info = gr.Textbox(label="Token counts (transcript / prompt / load path)", lines=2) gr.Markdown("
") with gr.Row(): with gr.Column(): with gr.Accordion("Instructions used (system prompt)", open=False): instr_md = gr.Markdown("```\n" + DEFAULT_SYSTEM_INSTRUCTIONS + "\n```") with gr.Column(): with gr.Accordion("Context used (glossary)", open=True): context_md = gr.Markdown("") # Reset labels reset_btn.click(fn=lambda: OFFICIAL_LABELS_TEXT, inputs=None, outputs=labels_text) # Warm-up warm_btn.click( fn=warmup_model, inputs=[repo, use_4bit, use_sdpa, hf_token, force_tok_redownload], outputs=diag ) def _pack_context_md(glossary_json, allowed_text): try: glossary = json.loads(glossary_json) if glossary_json else DEFAULT_LABEL_GLOSSARY except Exception: glossary = DEFAULT_LABEL_GLOSSARY allowed_list = [ln.strip() for ln in (allowed_text or OFFICIAL_LABELS_TEXT).splitlines() if ln.strip()] return "### Label Glossary (used)\n" + "\n".join(f"- {k}: {glossary.get(k,'')}" for k in allowed_list) context_md.value = _pack_context_md(json.dumps(DEFAULT_LABEL_GLOSSARY), OFFICIAL_LABELS_TEXT) # Run single run_btn.click( fn=run_single, inputs=[ text, file, gt_text, gt_file, use_cleaning, use_fallback, labels_text, sys_instr_tb, glossary_tb, fallback_tb, repo, use_4bit, use_sdpa, max_tokens, hf_token, force_tok_redownload ], outputs=[summary, json_out, diag, raw, context_md, instr_md, metrics_tb, prompt_preview, token_info], ) with gr.Tab("Batch evaluation"): with gr.Row(): with gr.Column(scale=3): gr.Markdown("
ZIP input
") zip_in = gr.File(label="ZIP with transcripts (.txt) and truths (.json)", file_types=[".zip"], type="filepath") use_cleaning_b = gr.Checkbox(label="Apply default cleaning", value=True) use_fallback_b = gr.Checkbox(label="Enable multilingual fallback rule layer", value=True) gr.Markdown("
") with gr.Column(scale=2): gr.Markdown("
Model & run
") repo_b = gr.Dropdown(label="Model", choices=MODEL_CHOICES, value=MODEL_CHOICES[0]) use_4bit_b = gr.Checkbox(label="Use 4-bit (GPU only)", value=True) use_sdpa_b = gr.Checkbox(label="Use SDPA attention (faster on many GPUs)", value=True) force_tok_redownload_b = gr.Checkbox(label="Force fresh tokenizer download", value=False) max_tokens_b = gr.Slider(label="Max input tokens", minimum=1024, maximum=8192, step=512, value=2048) hf_token_b = gr.Textbox(label="HF_TOKEN (only for gated models)", type="password", value=os.environ.get("HF_TOKEN","")) sys_instr_tb_b = gr.Textbox(label="System Instructions (editable for batch)", value=DEFAULT_SYSTEM_INSTRUCTIONS, lines=4) glossary_tb_b = gr.Code(label="Label Glossary (JSON; editable for batch)", value=json.dumps(DEFAULT_LABEL_GLOSSARY, indent=2), language="json") fallback_tb_b = gr.Code(label="Fallback Cues (Multilingual, JSON; editable for batch)", value=json.dumps(DEFAULT_FALLBACK_CUES, indent=2), language="json") limit_files = gr.Slider(label="Process at most N files (0 = all)", minimum=0, maximum=2000, step=10, value=0) run_batch_btn = gr.Button("Run Batch", variant="primary") gr.Markdown("
") with gr.Row(): gr.Markdown("
Batch outputs
") status = gr.Textbox(label="Status", lines=1) diag_b = gr.Textbox(label="Batch diagnostics & metrics", lines=12) df_out = gr.Dataframe(label="Per-file results (TP/FP/FN, latency)", interactive=False) csv_out = gr.File(label="Download CSV", interactive=False) gr.Markdown("
") run_batch_btn.click( fn=run_batch, inputs=[ zip_in, use_cleaning_b, use_fallback_b, sys_instr_tb_b, glossary_tb_b, fallback_tb_b, repo_b, use_4bit_b, use_sdpa_b, max_tokens_b, hf_token_b, force_tok_redownload_b, limit_files ], outputs=[status, diag_b, df_out, csv_out], ) if __name__ == "__main__": # Optional: print environment info to logs try: print("Torch version:", torch.__version__) print("CUDA available:", torch.cuda.is_available()) if torch.cuda.is_available(): print("CUDA (compiled):", torch.version.cuda) print("Device:", torch.cuda.get_device_name(0)) except Exception as _: pass demo.launch()