Spaces:
Sleeping
Sleeping
| import ast | |
| import json | |
| import logging | |
| import os | |
| import re | |
| import threading | |
| import time | |
| from pathlib import Path | |
| from typing import Any | |
| import unsloth | |
| import torch | |
| from langchain_community.embeddings import HuggingFaceEmbeddings | |
| from langchain_community.vectorstores import FAISS | |
| from unsloth import FastLanguageModel | |
| from config import ( | |
| EMBEDDING_MODEL_NAME, | |
| HF_TOKEN, | |
| JURISDICTIONS, | |
| LLM_MODEL_NAME, | |
| LOAD_IN_4BIT, | |
| MAX_ACCUMULATED_CHUNKS, | |
| MAX_CONTEXT_CHARACTERS, | |
| MAX_SEMANTIC_QUERY_CHARACTERS, | |
| MAX_SEMANTIC_QUERY_WORDS, | |
| MAX_SEQ_LENGTH, | |
| MODEL_DTYPE, | |
| jurisdiction_db_path, | |
| jurisdiction_name, | |
| normalize_jurisdiction, | |
| ) | |
| logger = logging.getLogger(__name__) | |
| _model = None | |
| _tokenizer = None | |
| _embeddings = None | |
| _vector_dbs: dict[str, FAISS] = {} | |
| _model_lock = threading.Lock() | |
| _embedding_lock = threading.Lock() | |
| _vector_lock = threading.Lock() | |
| _generation_lock = threading.Lock() | |
| _counter_lock = threading.Lock() | |
| _app_started_at = time.time() | |
| _requests_started = 0 | |
| _requests_completed = 0 | |
| _requests_failed = 0 | |
| def record_request_started() -> int: | |
| global _requests_started | |
| with _counter_lock: | |
| _requests_started += 1 | |
| return _requests_started | |
| def record_request_finished(success: bool) -> None: | |
| global _requests_completed, _requests_failed | |
| with _counter_lock: | |
| if success: | |
| _requests_completed += 1 | |
| else: | |
| _requests_failed += 1 | |
| def is_model_loaded() -> bool: | |
| return _model is not None and _tokenizer is not None | |
| def loaded_vector_db_codes() -> list[str]: | |
| with _vector_lock: | |
| return sorted(_vector_dbs.keys()) | |
| def is_vector_db_loaded(jurisdiction: str | None = None) -> bool: | |
| if jurisdiction is None: | |
| return bool(loaded_vector_db_codes()) | |
| code = normalize_jurisdiction(jurisdiction) | |
| with _vector_lock: | |
| return code in _vector_dbs | |
| def database_files_ready(jurisdiction: str) -> bool: | |
| db_path = jurisdiction_db_path(jurisdiction) | |
| return ( | |
| (db_path / "index.faiss").is_file() | |
| and (db_path / "index.pkl").is_file() | |
| ) | |
| def initialize_llm(): | |
| global _model, _tokenizer | |
| if is_model_loaded(): | |
| return _model, _tokenizer | |
| with _model_lock: | |
| if is_model_loaded(): | |
| return _model, _tokenizer | |
| if not torch.cuda.is_available(): | |
| raise RuntimeError( | |
| "CUDA is unavailable. This Space requires an NVIDIA GPU." | |
| ) | |
| kwargs = { | |
| "model_name": LLM_MODEL_NAME, | |
| "max_seq_length": MAX_SEQ_LENGTH, | |
| "dtype": MODEL_DTYPE, | |
| "load_in_4bit": LOAD_IN_4BIT, | |
| } | |
| if HF_TOKEN: | |
| kwargs["token"] = HF_TOKEN | |
| logger.info("Loading model: %s", LLM_MODEL_NAME) | |
| _model, _tokenizer = FastLanguageModel.from_pretrained(**kwargs) | |
| FastLanguageModel.for_inference(_model) | |
| _model.eval() | |
| if _tokenizer.pad_token_id is None: | |
| _tokenizer.pad_token = _tokenizer.eos_token | |
| logger.info("Model loaded.") | |
| return _model, _tokenizer | |
| def initialize_embeddings(): | |
| global _embeddings | |
| if _embeddings is not None: | |
| return _embeddings | |
| with _embedding_lock: | |
| if _embeddings is None: | |
| logger.info("Loading embedding model: %s", EMBEDDING_MODEL_NAME) | |
| _embeddings = HuggingFaceEmbeddings( | |
| model_name=EMBEDDING_MODEL_NAME | |
| ) | |
| return _embeddings | |
| def initialize_vector_db(jurisdiction: str): | |
| code = normalize_jurisdiction(jurisdiction) | |
| with _vector_lock: | |
| if code in _vector_dbs: | |
| return _vector_dbs[code] | |
| db_path = jurisdiction_db_path(code) | |
| if not database_files_ready(code): | |
| raise FileNotFoundError( | |
| f"The {jurisdiction_name(code)} FAISS database is incomplete. " | |
| f"Expected '{db_path / 'index.faiss'}' and " | |
| f"'{db_path / 'index.pkl'}'." | |
| ) | |
| embeddings = initialize_embeddings() | |
| logger.info( | |
| "Loading %s FAISS database: %s", | |
| jurisdiction_name(code), | |
| db_path, | |
| ) | |
| vector_db = FAISS.load_local( | |
| str(db_path), | |
| embeddings=embeddings, | |
| allow_dangerous_deserialization=True, | |
| ) | |
| _vector_dbs[code] = vector_db | |
| logger.info("%s FAISS database loaded.", jurisdiction_name(code)) | |
| return vector_db | |
| def call_llm(prompt: str, max_new_tokens: int) -> str: | |
| model, tokenizer = initialize_llm() | |
| chat = tokenizer.apply_chat_template( | |
| [{"role": "user", "content": prompt}], | |
| tokenize=False, | |
| add_generation_prompt=True, | |
| ) | |
| inputs = tokenizer( | |
| chat, | |
| return_tensors="pt", | |
| truncation=True, | |
| max_length=max(256, MAX_SEQ_LENGTH - max_new_tokens), | |
| ) | |
| device = next(model.parameters()).device | |
| inputs = {name: tensor.to(device) for name, tensor in inputs.items()} | |
| input_length = inputs["input_ids"].shape[-1] | |
| with _generation_lock, torch.inference_mode(): | |
| output = model.generate( | |
| **inputs, | |
| max_new_tokens=max_new_tokens, | |
| do_sample=False, | |
| use_cache=True, | |
| pad_token_id=tokenizer.pad_token_id, | |
| eos_token_id=tokenizer.eos_token_id, | |
| ) | |
| return tokenizer.decode( | |
| output[0, input_length:], | |
| skip_special_tokens=True, | |
| ).strip() | |
| def normalize_metadata(metadata: dict | None) -> dict: | |
| metadata = dict(metadata or {}) | |
| source = str(metadata.get("source") or "Unknown source") | |
| page = metadata.get("page") | |
| if page in (None, "", 0, "0"): | |
| page = "Unavailable" | |
| page_reference_basis = str( | |
| metadata.get("page_reference_basis") or "" | |
| ).strip() | |
| page_label = str( | |
| metadata.get("page_label") or "" | |
| ).strip() | |
| if not page_label: | |
| if str(page).lower() in { | |
| "unavailable", | |
| "unknown", | |
| "none", | |
| "null", | |
| "", | |
| }: | |
| page_label = "Page unavailable" | |
| elif page_reference_basis == "normalized_pdf": | |
| page_label = f"Page {page} (normalized PDF)" | |
| else: | |
| page_label = f"Page {page}" | |
| return { | |
| "title": str( | |
| metadata.get("title") | |
| or Path(source).stem.replace("_", " ").strip() | |
| or "Untitled document" | |
| ), | |
| "source": source, | |
| "url": str(metadata.get("url") or "").strip(), | |
| "page": page, | |
| "page_label": page_label, | |
| "page_reference_basis": page_reference_basis, | |
| "start_index": metadata.get("start_index"), | |
| "original_extension": str( | |
| metadata.get("original_extension") | |
| or Path(source).suffix.lower() | |
| ), | |
| "processing_extension": str( | |
| metadata.get("processing_extension") or "" | |
| ), | |
| "document_type": str( | |
| metadata.get("document_type") or "" | |
| ), | |
| } | |
| def retrieve_documents( | |
| query: str, | |
| top_k: int, | |
| jurisdiction: str, | |
| ) -> list[dict]: | |
| code = normalize_jurisdiction(jurisdiction) | |
| results = initialize_vector_db(code).similarity_search_with_score( | |
| query, | |
| k=top_k, | |
| ) | |
| chunks = [] | |
| for document, score in results: | |
| text = str(document.page_content or "").strip() | |
| if text: | |
| chunks.append( | |
| { | |
| "text": text, | |
| "metadata": normalize_metadata(document.metadata), | |
| "score": float(score), | |
| "jurisdiction": code, | |
| } | |
| ) | |
| return chunks | |
| def chunk_key(chunk: dict) -> tuple[str, str, str, str, str]: | |
| metadata = chunk["metadata"] | |
| return ( | |
| str(chunk.get("jurisdiction", "")), | |
| str(metadata.get("source", "")), | |
| str(metadata.get("page", "")), | |
| str(metadata.get("start_index", "")), | |
| chunk["text"].strip(), | |
| ) | |
| def add_unique_chunks( | |
| accumulated: list[dict], | |
| new_chunks: list[dict], | |
| seen: set, | |
| ) -> int: | |
| added = 0 | |
| for chunk in new_chunks: | |
| key = chunk_key(chunk) | |
| if key in seen: | |
| continue | |
| seen.add(key) | |
| accumulated.append(chunk) | |
| added += 1 | |
| if len(accumulated) >= MAX_ACCUMULATED_CHUNKS: | |
| break | |
| return added | |
| def page_display(metadata: dict) -> str: | |
| """Return the final human-readable page label.""" | |
| page_label = str( | |
| metadata.get("page_label") or "" | |
| ).strip() | |
| if page_label: | |
| return page_label | |
| page = metadata.get("page") | |
| if page in (None, "", 0, "0"): | |
| return "Page unavailable" | |
| if str(metadata.get("page_reference_basis") or "") == "normalized_pdf": | |
| return f"Page {page} (normalized PDF)" | |
| return f"Page {page}" | |
| def build_context(chunks: list[dict]) -> str: | |
| sections = [] | |
| used = 0 | |
| for number, chunk in enumerate(chunks, 1): | |
| metadata = chunk["metadata"] | |
| jurisdiction_code = str(chunk.get("jurisdiction", "")) | |
| jurisdiction_label = ( | |
| jurisdiction_name(jurisdiction_code) | |
| if jurisdiction_code in JURISDICTIONS | |
| else "" | |
| ) | |
| document_name = str( | |
| chunk.get("collection_label") | |
| or metadata.get("title") | |
| or Path(str(metadata.get("source") or "")).stem.replace("_", " ").strip() | |
| or "Uploaded document" | |
| ) | |
| metadata_lines = [ | |
| f"[Source {number}]", | |
| f"Document: {document_name}", | |
| f"Title: {metadata['title']}", | |
| f"Location: {page_display(metadata)}", | |
| ] | |
| if jurisdiction_label: | |
| metadata_lines.append(f"Jurisdiction: {jurisdiction_label}") | |
| section = ( | |
| "\n".join(metadata_lines) | |
| + "\nRetrieved text:\n" | |
| + chunk["text"] | |
| ) | |
| if sections and used + len(section) > MAX_CONTEXT_CHARACTERS: | |
| break | |
| sections.append(section) | |
| used += len(section) | |
| return "\n\n".join(sections) | |
| def format_sources(chunks: list[dict]) -> str: | |
| """Return predictable Markdown that the public frontend can parse.""" | |
| if not chunks: | |
| return "No source passages were retrieved." | |
| cards = [] | |
| for number, chunk in enumerate(chunks, 1): | |
| metadata = chunk["metadata"] | |
| preview = chunk["text"].strip() | |
| if len(preview) > 1200: | |
| preview = preview[:1200].rstrip() + "…" | |
| jurisdiction_code = str(chunk.get("jurisdiction", "")) | |
| jurisdiction_label = ( | |
| jurisdiction_name(jurisdiction_code) | |
| if jurisdiction_code in JURISDICTIONS | |
| else "" | |
| ) | |
| document_name = str( | |
| chunk.get("collection_label") | |
| or metadata.get("title") | |
| or Path(str(metadata.get("source") or "")).stem.replace("_", " ").strip() | |
| or "Uploaded document" | |
| ) | |
| title = str(metadata.get("title") or Path(document_name).name) | |
| metadata_parts = [ | |
| f"**Source:** {number}", | |
| f"**Page:** {page_display(metadata)}", | |
| f"**Document:** `{document_name}`", | |
| f"**Retrieval distance:** {chunk['score']:.4f}", | |
| ] | |
| if jurisdiction_label: | |
| metadata_parts.insert( | |
| 2, | |
| f"**Jurisdiction:** {jurisdiction_label}", | |
| ) | |
| url = str(metadata.get("url") or "").strip() | |
| if url: | |
| metadata_parts.append(f"**URL:** `{url}`") | |
| cards.append( | |
| f"### [{number}] {title}\n" | |
| + " | ".join(metadata_parts) | |
| + "\n\n" | |
| + preview | |
| ) | |
| return "\n\n---\n\n".join(cards) | |
| def normalize_source_citations( | |
| answer: str, | |
| source_count: int, | |
| ) -> str: | |
| """Keep answer citations aligned with the visible source cards.""" | |
| text = str(answer or "") | |
| if source_count < 1: | |
| return re.sub(r"\[[0-9,\s-]+\]", "", text) | |
| def replace(match: re.Match) -> str: | |
| raw = match.group(1).strip() | |
| if "," in raw or " " in raw: | |
| values = [int(value) for value in re.findall(r"\d+", raw)] | |
| valid = [ | |
| value for value in values | |
| if 1 <= value <= source_count | |
| ] | |
| return "".join(f"[{value}]" for value in valid) | |
| try: | |
| number = int(raw) | |
| except ValueError: | |
| return "" | |
| if 1 <= number <= source_count: | |
| return f"[{number}]" | |
| digits = [int(char) for char in raw if char.isdigit()] | |
| if ( | |
| len(digits) > 1 | |
| and all(1 <= value <= source_count for value in digits) | |
| ): | |
| return "".join(f"[{value}]" for value in digits) | |
| return "" | |
| text = re.sub(r"\[([0-9,\s]+)\]", replace, text) | |
| return re.sub(r"[ \t]{2,}", " ", text).strip() | |
| def compact_semantic_query(text: str) -> str: | |
| """Normalize and cap a semantic-search query without domain hard-coding.""" | |
| query = re.sub(r"\s+", " ", str(text or "")).strip().strip("\"'") | |
| query = re.sub(r"\b(?:SELECT|FROM|WHERE)\b", " ", query, flags=re.IGNORECASE) | |
| query = re.sub(r"\s+", " ", query).strip(" ,;|") | |
| words = query.split() | |
| if len(words) > MAX_SEMANTIC_QUERY_WORDS: | |
| query = " ".join(words[:MAX_SEMANTIC_QUERY_WORDS]) | |
| if len(query) > MAX_SEMANTIC_QUERY_CHARACTERS: | |
| query = query[:MAX_SEMANTIC_QUERY_CHARACTERS].rsplit(" ", 1)[0] | |
| return query.strip() | |
| def clean_search_query(text: str) -> str: | |
| query = re.sub( | |
| r"^```.*?\n|\n```$", | |
| "", | |
| str(text or "").strip(), | |
| flags=re.DOTALL, | |
| ) | |
| return compact_semantic_query(query) | |
| def _extract_balanced_json_object(text: str) -> str | None: | |
| start = text.find("{") | |
| if start < 0: | |
| return None | |
| depth = 0 | |
| in_string = False | |
| escaped = False | |
| for index in range(start, len(text)): | |
| char = text[index] | |
| if in_string: | |
| if escaped: | |
| escaped = False | |
| elif char == "\\": | |
| escaped = True | |
| elif char == '"': | |
| in_string = False | |
| continue | |
| if char == '"': | |
| in_string = True | |
| elif char == "{": | |
| depth += 1 | |
| elif char == "}": | |
| depth -= 1 | |
| if depth == 0: | |
| return text[start:index + 1] | |
| return None | |
| def _fallback_judgment_fields(text: str) -> dict: | |
| result: dict[str, Any] = {} | |
| boolean = re.search( | |
| r'["\']?sufficient["\']?\s*[:=]\s*(true|false|yes|no)', | |
| text, | |
| flags=re.IGNORECASE, | |
| ) | |
| if boolean: | |
| result["sufficient"] = boolean.group(1).lower() in {"true", "yes"} | |
| for key in ("missing_information", "reason", "next_search_focus"): | |
| match = re.search( | |
| rf'["\']?{key}["\']?\s*[:=]\s*["\']?(.*?)(?=\n\s*["\']?(?:sufficient|missing_information|reason|next_search_focus|cross_references)["\']?\s*[:=]|$)', | |
| text, | |
| flags=re.IGNORECASE | re.DOTALL, | |
| ) | |
| if match: | |
| result[key] = match.group(1).strip().strip('"\' ,') | |
| refs = re.search( | |
| r'["\']?cross_references["\']?\s*[:=]\s*(\[[^\]]*\]|[^\n]+)', | |
| text, | |
| flags=re.IGNORECASE, | |
| ) | |
| if refs: | |
| result["cross_references"] = normalize_regulatory_references(refs.group(1)) | |
| return result | |
| def parse_judgment(text: str) -> dict: | |
| """Parse structured model output with tolerant, deterministic fallbacks.""" | |
| cleaned = re.sub( | |
| r"```(?:json)?|```", | |
| "", | |
| str(text or ""), | |
| flags=re.IGNORECASE, | |
| ).strip() | |
| candidates = [cleaned] | |
| balanced = _extract_balanced_json_object(cleaned) | |
| if balanced and balanced != cleaned: | |
| candidates.append(balanced) | |
| for candidate in candidates: | |
| try: | |
| value = json.loads(candidate) | |
| if isinstance(value, dict): | |
| return value | |
| except (json.JSONDecodeError, TypeError): | |
| pass | |
| try: | |
| pythonish = re.sub(r"\btrue\b", "True", candidate, flags=re.IGNORECASE) | |
| pythonish = re.sub(r"\bfalse\b", "False", pythonish, flags=re.IGNORECASE) | |
| pythonish = re.sub(r"\bnull\b", "None", pythonish, flags=re.IGNORECASE) | |
| value = ast.literal_eval(pythonish) | |
| if isinstance(value, dict): | |
| return value | |
| except (ValueError, SyntaxError): | |
| pass | |
| fallback = _fallback_judgment_fields(cleaned) | |
| if "sufficient" in fallback: | |
| return fallback | |
| raise ValueError("No usable judgment object was found.") | |
| _FULL_RULE_REFERENCE = re.compile( | |
| r"\b\d{1,3}[A-Za-z]?[-–]\d{1,3}(?:\.\d+)+(?:\([A-Za-z0-9]+\))*(?=\W|$)", | |
| re.IGNORECASE, | |
| ) | |
| _STATUTE_REFERENCE = re.compile( | |
| r"(?<![-\d])\b\d{3}\.\d+(?:\([A-Za-z0-9]+\))*(?=\W|$)", | |
| re.IGNORECASE, | |
| ) | |
| _TABLE_APPENDIX_REFERENCE = re.compile( | |
| r"\b(?:Table|Appendix)\s+[A-Za-z0-9.-]+\b", | |
| re.IGNORECASE, | |
| ) | |
| def canonicalize_regulatory_reference(value: str) -> str: | |
| text = re.sub(r"\s+", " ", str(value or "")).strip(" .,;:\"'") | |
| text = re.sub( | |
| r"^(?:Rule|Chapter|section|subsection|paragraph|subparagraph)\s+", | |
| "", | |
| text, | |
| flags=re.IGNORECASE, | |
| ) | |
| text = text.replace("–", "-") | |
| full = _FULL_RULE_REFERENCE.search(text) | |
| if full: | |
| return full.group(0) | |
| table = _TABLE_APPENDIX_REFERENCE.search(text) | |
| if table: | |
| prefix, suffix = table.group(0).split(" ", 1) | |
| return f"{prefix.title()} {suffix}" | |
| statute = _STATUTE_REFERENCE.search(text) | |
| if statute: | |
| return statute.group(0) | |
| return "" | |
| def normalize_regulatory_references(value: Any, limit: int = 12) -> list[str]: | |
| """Return unique canonical citations and exclude paths or prose.""" | |
| if value is None: | |
| return [] | |
| if isinstance(value, (list, tuple, set)): | |
| raw_items = [str(item or "") for item in value] | |
| else: | |
| raw_items = [str(value or "")] | |
| candidates: list[str] = [] | |
| for raw in raw_items: | |
| candidates.extend(match.group(0) for match in _FULL_RULE_REFERENCE.finditer(raw)) | |
| candidates.extend(match.group(0) for match in _TABLE_APPENDIX_REFERENCE.finditer(raw)) | |
| occupied = [match.span() for match in _FULL_RULE_REFERENCE.finditer(raw)] | |
| for match in _STATUTE_REFERENCE.finditer(raw): | |
| if not any(start <= match.start() < end for start, end in occupied): | |
| candidates.append(match.group(0)) | |
| canonical_candidates = [ | |
| canonicalize_regulatory_reference(candidate) | |
| for candidate in candidates | |
| ] | |
| full_rule_suffixes = { | |
| item.split("-", 1)[1].casefold() | |
| for item in canonical_candidates | |
| if item and "-" in item | |
| } | |
| result: list[str] = [] | |
| seen: set[str] = set() | |
| for canonical in canonical_candidates: | |
| key = canonical.casefold() | |
| if not canonical: | |
| continue | |
| if "-" not in canonical and key in full_rule_suffixes: | |
| continue | |
| if key not in seen: | |
| seen.add(key) | |
| result.append(canonical) | |
| if len(result) >= limit: | |
| break | |
| return result | |
| def extract_regulatory_cross_references(chunks: list[dict]) -> list[str]: | |
| """Extract and canonicalize rule, section, table, and appendix leads.""" | |
| return normalize_regulatory_references( | |
| [str(chunk.get("text") or "") for chunk in chunks], | |
| limit=12, | |
| ) | |
| def normalize_string_list(value: Any) -> list[str]: | |
| """Normalize a model-produced string/list field to unique strings.""" | |
| if value is None: | |
| return [] | |
| if isinstance(value, str): | |
| raw_values = re.split(r"[\n,;]+", value) | |
| elif isinstance(value, (list, tuple, set)): | |
| raw_values = list(value) | |
| else: | |
| raw_values = [value] | |
| result: list[str] = [] | |
| seen: set[str] = set() | |
| for item in raw_values: | |
| text = re.sub(r"\s+", " ", str(item or "")).strip(" -\t") | |
| key = text.casefold() | |
| if text and key not in seen: | |
| seen.add(key) | |
| result.append(text) | |
| return result | |
| def looks_like_sql_query(query: str) -> bool: | |
| text = str(query or "") | |
| return bool(re.search(r"\b(?:select|from|where)\b|\w+\s*=", text, re.IGNORECASE)) | |
| def semantic_query_fallback( | |
| question: str, | |
| target_concept: str = "", | |
| cross_references: list[str] | None = None, | |
| next_search_focus: str = "", | |
| ) -> str: | |
| """Build a compact non-SQL query when model query generation fails.""" | |
| parts: list[str] = [] | |
| for value in [ | |
| *(normalize_regulatory_references(cross_references or [])[:2]), | |
| target_concept, | |
| next_search_focus, | |
| question, | |
| ]: | |
| text = re.sub(r"\s+", " ", str(value or "")).strip() | |
| if text and text.casefold() not in {part.casefold() for part in parts}: | |
| parts.append(text) | |
| return compact_semantic_query(" ".join(parts)) | |
| def build_exact_reference_query( | |
| reference: str, | |
| target_concept: str, | |
| question_type: str, | |
| ) -> str: | |
| """Build one compact citation-focused FAISS query.""" | |
| canonical = canonicalize_regulatory_reference(reference) | |
| evidence_hint = { | |
| "definition": "definition means defined as", | |
| "numeric_criterion": "value units criterion", | |
| "calculation": "formula calculation", | |
| "requirement": "requirement reduction allocation", | |
| "applicability": "applicability applies", | |
| "procedure": "procedure steps", | |
| "exception": "exception conditions", | |
| }.get(str(question_type or "").strip().lower(), "") | |
| return compact_semantic_query( | |
| " ".join(part for part in [canonical, target_concept, evidence_hint] if part) | |
| ) | |
| def has_explicit_definition_evidence( | |
| chunks: list[dict], | |
| target_concept: str, | |
| ) -> bool: | |
| """Check that a definition claim is backed by actual defining language.""" | |
| target = re.sub(r"\s+", " ", str(target_concept or "")).strip() | |
| if not target: | |
| return False | |
| variants = [target] | |
| shortened = re.sub( | |
| r"\s+(?:waters?|standards?|criteria|term|definition)$", | |
| "", | |
| target, | |
| flags=re.IGNORECASE, | |
| ).strip() | |
| if shortened and shortened.casefold() != target.casefold(): | |
| variants.append(shortened) | |
| patterns: list[str] = [] | |
| for variant in variants: | |
| escaped = re.escape(variant) | |
| patterns.extend([ | |
| rf'["“]?{escaped}["”]?\s+(?:shall\s+mean|means|is\s+defined\s+as|refers\s+to)\b', | |
| rf'(?:definition\s+of|define(?:s|d)?\s+)["“]?{escaped}["”]?\s*(?:is|as|:)\b', | |
| rf'["“]?{escaped}["”]?\s*[:—-]\s+\S+', | |
| ]) | |
| for chunk in chunks: | |
| text = re.sub(r"\s+", " ", str(chunk.get("text") or "")) | |
| for pattern in patterns: | |
| if re.search(pattern, text, flags=re.IGNORECASE): | |
| return True | |
| return False | |
| def format_uptime(seconds: float) -> str: | |
| seconds = max(0, int(seconds)) | |
| days, seconds = divmod(seconds, 86400) | |
| hours, seconds = divmod(seconds, 3600) | |
| minutes, seconds = divmod(seconds, 60) | |
| parts = [] | |
| if days: | |
| parts.append(f"{days}d") | |
| if hours or days: | |
| parts.append(f"{hours}h") | |
| if minutes or hours or days: | |
| parts.append(f"{minutes}m") | |
| parts.append(f"{seconds}s") | |
| return " ".join(parts) | |
| def gpu_status_lines() -> list[str]: | |
| if not torch.cuda.is_available(): | |
| return ["CUDA available: No"] | |
| device = torch.cuda.current_device() | |
| properties = torch.cuda.get_device_properties(device) | |
| gib = 1024**3 | |
| return [ | |
| "CUDA available: Yes", | |
| f"GPU: {properties.name}", | |
| f"GPU memory total: {properties.total_memory / gib:.2f} GB", | |
| f"GPU memory allocated: {torch.cuda.memory_allocated(device) / gib:.2f} GB", | |
| f"GPU memory reserved: {torch.cuda.memory_reserved(device) / gib:.2f} GB", | |
| ] | |
| def check_system_status() -> str: | |
| with _counter_lock: | |
| counters = ( | |
| _requests_started, | |
| _requests_completed, | |
| _requests_failed, | |
| ) | |
| loaded_codes = loaded_vector_db_codes() | |
| loaded_names = [jurisdiction_name(code) for code in loaded_codes] | |
| lines = [ | |
| f"Process ID: {os.getpid()}", | |
| f"Process uptime: {format_uptime(time.time() - _app_started_at)}", | |
| "Available jurisdictions: " | |
| + ", ".join( | |
| f"{details['name']} ({code.upper()})" | |
| for code, details in JURISDICTIONS.items() | |
| ), | |
| ] | |
| for code, details in JURISDICTIONS.items(): | |
| db_path = jurisdiction_db_path(code) | |
| lines.extend( | |
| [ | |
| f"{details['name']} database files: " | |
| + ("Ready" if database_files_ready(code) else "Missing index.faiss and/or index.pkl"), | |
| f"{details['name']} database path: {db_path}", | |
| f"{details['name']} vector database loaded in memory: " | |
| + ("Yes" if code in loaded_codes else "No"), | |
| ] | |
| ) | |
| lines.extend( | |
| [ | |
| "Loaded vector databases: " | |
| + (", ".join(loaded_names) if loaded_names else "None"), | |
| "LLM provider: Direct local Unsloth inference", | |
| f"LLM model: {LLM_MODEL_NAME}", | |
| f"Model loaded in memory: {'Yes' if is_model_loaded() else 'No'}", | |
| f"Model dtype: {str(MODEL_DTYPE).replace('torch.', '')}", | |
| f"4-bit loading: {'Yes' if LOAD_IN_4BIT else 'No'}", | |
| f"Embedding model: {EMBEDDING_MODEL_NAME}", | |
| f"Requests started: {counters[0]}", | |
| f"Requests completed: {counters[1]}", | |
| f"Requests failed/empty: {counters[2]}", | |
| ] | |
| ) | |
| return "\n".join(lines + gpu_status_lines()) | |