import torch import numpy as np import threading from contextlib import asynccontextmanager from fastapi import FastAPI from pydantic import BaseModel from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModel, AutoModelForSeq2SeqLM from sentence_transformers import SentenceTransformer from sklearn.metrics.pairwise import cosine_similarity from setfit import SetFitModel from gliner import GLiNER from typing import List, Optional import os os.environ["TOKENIZERS_PARALLELISM"] = "false" models = {} hf_lock = threading.Lock() MAX_TEXT_CHARS = 50_000 # ---------- TextChunker (from Raubachm/sentence-transformers-semantic-chunker) ---------- class TextChunker: def __init__(self, st_model: SentenceTransformer): self.model = st_model def chunk(self, text: str, context_window: int = 1, percentile_threshold: float = 75, min_chunk_size: int = 2,max_chunk_tokens: int = 400) -> List[str]: import nltk nltk.download("punkt", quiet=True) nltk.download("punkt_tab", quiet=True) from nltk.tokenize import sent_tokenize sentences = sent_tokenize(text) if len(sentences) <= 1: return [text] if text.strip() else [] contextualized = self._add_context(sentences, context_window) embeddings = self.model.encode(contextualized,batch_size=32,show_progress_bar=False) distances = self._calculate_distances(embeddings) if not distances: return [text] effective_threshold = percentile_threshold if len(sentences) < 20: effective_threshold = min(percentile_threshold, 70.0) breakpoints = self._identify_breakpoints(distances, effective_threshold) initial_chunks = self._create_chunks(sentences, breakpoints) chunk_embeddings = self.model.encode(initial_chunks,batch_size=32,show_progress_bar=False) merged_chunks = self._merge_small_chunks(initial_chunks, chunk_embeddings, min_chunk_size) final_chunks = [] for chunk in merged_chunks: if self._estimate_tokens(chunk) > max_chunk_tokens: sub_chunks = self._split_oversized(chunk, max_chunk_tokens, sent_tokenize) final_chunks.extend(sub_chunks) else: final_chunks.append(chunk) return [c for c in final_chunks if c.strip()] def _estimate_tokens(self, text: str) -> int: # Fast approximation: 1 token ≈ 4 chars for English legal text return len(text) // 4 def _split_oversized(self, chunk: str, max_tokens: int, sent_tokenize) -> List[str]: """Split a chunk that exceeds max_tokens at sentence boundaries.""" sentences = sent_tokenize(chunk) result, current, current_tokens = [], [], 0 for sent in sentences: t = self._estimate_tokens(sent) if current_tokens + t > max_tokens and current: result.append(' '.join(current)) current, current_tokens = [sent], t else: current.append(sent) current_tokens += t if current: result.append(' '.join(current)) return result def _add_context(self, sentences, window_size): result = [] for i in range(len(sentences)): start = max(0, i - window_size) end = min(len(sentences), i + window_size + 1) result.append(" ".join(sentences[start:end])) return result def _calculate_distances(self, embeddings): distances = [] for i in range(len(embeddings) - 1): sim = cosine_similarity([embeddings[i]], [embeddings[i + 1]])[0][0] distances.append(1 - sim) return distances def _identify_breakpoints(self, distances, threshold_percentile): threshold = np.percentile(distances, threshold_percentile) return [i for i, d in enumerate(distances) if d > threshold] def _create_chunks(self, sentences, breakpoints): chunks, start = [], 0 for bp in breakpoints: chunks.append(" ".join(sentences[start:bp + 1])) start = bp + 1 chunks.append(" ".join(sentences[start:])) return chunks def _merge_small_chunks(self, chunks, embeddings, min_size): if len(chunks) <= 1: return chunks final_chunks = [chunks[0]] merged_embeddings = [embeddings[0]] for i in range(1, len(chunks) - 1): if len(chunks[i].split(". ")) < min_size: prev_sim = cosine_similarity([embeddings[i]], [merged_embeddings[-1]])[0][0] next_sim = cosine_similarity([embeddings[i]], [embeddings[i + 1]])[0][0] if prev_sim > next_sim: final_chunks[-1] = f"{final_chunks[-1]} {chunks[i]}" merged_embeddings[-1] = (merged_embeddings[-1] + embeddings[i]) / 2 else: chunks[i + 1] = f"{chunks[i]} {chunks[i + 1]}" embeddings[i + 1] = (embeddings[i] + embeddings[i + 1]) / 2 else: final_chunks.append(chunks[i]) merged_embeddings.append(embeddings[i]) final_chunks.append(chunks[-1]) return final_chunks # ---------- Lifespan ---------- @asynccontextmanager async def lifespan(app: FastAPI): print("Loading models...") # Helper to load quantized BERT models with optimum-quanto def load_quantized_bert(model_id, num_labels=None): from transformers import QuantoConfig quant_config = QuantoConfig(weights="int8") tokenizer = AutoTokenizer.from_pretrained(model_id) model = AutoModelForSequenceClassification.from_pretrained( model_id, num_labels=num_labels, quantization_config=quant_config, # will be ignored for AutoModel (num_labels=None) ignore_mismatched_sizes=True ) if num_labels else AutoModel.from_pretrained( model_id, quantization_config=quant_config ) model.eval() return model, tokenizer # 1. SetFit contracts clauses print("Loading SetFit contracts clauses model...") models["contracts_clauses"] = SetFitModel.from_pretrained( "scholarly360/setfit-contracts-clauses" ) print("✓ contracts_clauses loaded") # 2. Contract NLI print("Loading contract NLI model(int8)...") nli_model,nli_tokenizer=load_quantized_bert("Syamchand/contract-nli-bert", num_labels=3) models["nli_tokenizer"] = nli_tokenizer models["nli_model"] = nli_model models["nli_id2label"] = {0: "entailment", 1: "neutral", 2: "contradiction"} print("✓ contract-nli loaded (int8) ") # 3. Clause risk classifier print("Loading clause risk classifier(int8)...") risk_model, risk_tokenizer = load_quantized_bert("Syamchand/clause_risk_classifier", num_labels=3) models["risk_tokenizer"] = risk_tokenizer models["risk_model"] = risk_model #models["risk_model"].eval() models["risk_id2label"] = {0: "low", 1: "medium", 2: "high"} print("✓ clause_risk_classifier loaded (int8)") # 5. Semantic chunker — load the backbone model specified in the Raubachm model card print("Loading semantic chunker model...") st_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", device="cpu") models["chunker"] = TextChunker(st_model) models["minilm_embeddings"] = st_model print("✓ semantic chunker loaded") # 6. NuNER_Zero NER model print("Loading NuNER_Zero NER model(int8)...") models["ner"] = GLiNER.from_pretrained("numind/NuNER_Zero",quantize=True) print("✓ NuNER_Zero loaded (int8)") print("All models ready!") yield models.clear() app = FastAPI(lifespan=lifespan) # ---------- Schemas ---------- class TextRequest(BaseModel): text: str class PairRequest(BaseModel): premise: str hypothesis: str class ChunkRequest(BaseModel): text: str percentile_threshold: float = 75.0 context_window: int = 1 min_chunk_size: int = 2 class ClassificationResult(BaseModel): label: str score: float class EmbeddingResult(BaseModel): embeddings: List[List[float]] class ChunkResult(BaseModel): chunks: List[str] class NERRequest(BaseModel): text: str entity_types: Optional[List[str]] = None class Entity(BaseModel): text: str label: str score: float start: int end: int class NERResult(BaseModel): entities: List[Entity] class BatchTextRequest(BaseModel): texts: List[str] class BatchClassificationResult(BaseModel): results: List[ClassificationResult] class RetrievalNLIRequest(BaseModel): """NLI with retrieval: finds most relevant clause chunks before inference.""" clauses: List[str] # Pre-segmented clauses (from segmenter) hypothesis: str top_k: int = 3 # Number of top clauses to concatenate as premise class RetrievalNLIResult(BaseModel): label: str score: float supporting_clauses: List[str] supporting_indices: List[int] class BatchNLIRequest(BaseModel): """Run multiple hypotheses against the same set of clauses.""" clauses: List[str] hypotheses: List[str] top_k: int = 3 class BatchNLIResult(BaseModel): results: List[RetrievalNLIResult] class BatchNERRequest(BaseModel): texts: List[str] entity_types: Optional[List[str]] = None class BatchNERResult(BaseModel): results: List[List[Entity]] class EmbeddingRequest(BaseModel): texts: List[str] class MiniLMEmbeddingResult(BaseModel): embeddings: List[List[float]] dimension: int model: str # ---------- Endpoints ---------- @app.get("/health") def health(): return {"status": "ok",} @app.get("/memory") def container_memory(): # Try cgroup v2 first (most common on HF Spaces) if os.path.exists("/sys/fs/cgroup/memory.current"): with open("/sys/fs/cgroup/memory.current") as f: usage = int(f.read().strip()) with open("/sys/fs/cgroup/memory.max") as f: limit_str = f.read().strip() limit = int(limit_str) if limit_str != "max" else None # Fallback to cgroup v1 elif os.path.exists("/sys/fs/cgroup/memory/memory.usage_in_bytes"): with open("/sys/fs/cgroup/memory/memory.usage_in_bytes") as f: usage = int(f.read().strip()) with open("/sys/fs/cgroup/memory/memory.limit_in_bytes") as f: limit = int(f.read().strip()) else: return {"error": "Cannot read container memory"} if limit is None: return {"usage_mb": round(usage / (1024*1024), 2), "limit_mb": "unlimited", "percent": "unknown"} return { "usage_mb": round(usage / (1024*1024), 2), "limit_mb": round(limit / (1024*1024), 2), "percent": round(usage / limit * 100, 1) } @app.post("/predict/contracts_clauses", response_model=ClassificationResult) def predict_contracts_clauses(req: TextRequest): model = models["contracts_clauses"] # The SetFit model predicts labels directly (no integer conversion needed) preds = model.predict([req.text]) label = preds[0] # Already a string like 'terms' # Try to get a confidence score using predict_proba if available score = 1.0 if hasattr(model, "predict_proba"): try: probs = model.predict_proba([req.text])[0] # model.labels stores the label strings in the order expected by predict_proba if hasattr(model, "labels") and model.labels is not None: # Find the index of the predicted label if label in model.labels: idx = model.labels.index(label) score = probs[idx] else: score = max(probs) else: score = max(probs) except Exception: score = 1.0 return ClassificationResult(label=label, score=round(float(score), 4)) @app.post("/predict/nli", response_model=ClassificationResult) def predict_nli(req: PairRequest): inputs = models["nli_tokenizer"]( req.premise, req.hypothesis, return_tensors="pt", truncation=True ) with torch.no_grad(): logits = models["nli_model"](**inputs).logits probs = torch.nn.functional.softmax(logits, dim=-1) class_id = torch.argmax(probs, dim=-1).item() return ClassificationResult( label=models["nli_id2label"][class_id], score=round(probs[0][class_id].item(), 4) ) @app.post("/predict/risk", response_model=ClassificationResult) def predict_risk(req: TextRequest): inputs = models["risk_tokenizer"]( req.text, return_tensors="pt", truncation=True, max_length=512 ) with torch.no_grad(): logits = models["risk_model"](**inputs).logits probs = torch.nn.functional.softmax(logits, dim=-1) class_id = torch.argmax(probs, dim=-1).item() return ClassificationResult( label=models["risk_id2label"][class_id], score=round(probs[0][class_id].item(), 4) ) @app.post("/predict/semantic_chunks", response_model=ChunkResult) def semantic_chunking(req: ChunkRequest): chunks = models["chunker"].chunk( text=req.text, context_window=req.context_window, percentile_threshold=req.percentile_threshold, min_chunk_size=req.min_chunk_size ) return ChunkResult(chunks=chunks) @app.post("/predict/ner", response_model=NERResult) def predict_ner(req: NERRequest): # Default entity types suitable for freelancer contracts default_types = [ "deposit percentage", "payment days", "cancellation fee percentage", "liability limit amount", "governing state", "notice period days", "confidentiality duration", "party name", "contract value", "effective date", ] labels = req.entity_types if req.entity_types else default_types # GLiNER expects lowercase labels for optimal performance labels = [l.lower() for l in labels] raw_entities = models["ner"].predict_entities(req.text, labels) return NERResult(entities=[Entity(**ent) for ent in raw_entities]) @app.post("/predict/contracts_clauses_batch", response_model=BatchClassificationResult) def predict_contracts_clauses_batch(req: BatchTextRequest): """Batch classify multiple clauses in one call. Critical for performance.""" model = models["contracts_clauses"] # SetFit handles batches natively preds = model.predict(req.texts) results = [] if hasattr(model, "predict_proba"): try: all_probs = model.predict_proba(req.texts) for label, probs in zip(preds, all_probs): score = 1.0 if hasattr(model, "labels") and model.labels and label in model.labels: score = float(probs[model.labels.index(label)]) else: score = float(max(probs)) results.append(ClassificationResult(label=str(label), score=round(score, 4))) return BatchClassificationResult(results=results) except Exception: pass return BatchClassificationResult(results=[ ClassificationResult(label=str(label), score=1.0) for label in preds ]) @app.post("/predict/risk_batch", response_model=BatchClassificationResult) def predict_risk_batch(req: BatchTextRequest): """Batch risk classification. Processes all clauses in one forward pass.""" tokenizer = models["risk_tokenizer"] model = models["risk_model"] id2label = models["risk_id2label"] results = [] # Process in sub-batches of 8 to avoid OOM BATCH_SIZE = 8 for i in range(0, len(req.texts), BATCH_SIZE): batch = req.texts[i:i+BATCH_SIZE] inputs = tokenizer(batch, return_tensors="pt", truncation=True, max_length=512, padding=True) with torch.no_grad(): logits = model(**inputs).logits probs = torch.nn.functional.softmax(logits, dim=-1) class_ids = torch.argmax(probs, dim=-1).tolist() for j, class_id in enumerate(class_ids): results.append(ClassificationResult( label=id2label[class_id], score=round(probs[j][class_id].item(), 4) )) return BatchClassificationResult(results=results) @app.post("/predict/nli_retrieval", response_model=RetrievalNLIResult) def predict_nli_retrieval(req: RetrievalNLIRequest): """ Retrieval-augmented NLI: 1. Embed all clauses + hypothesis using the sentence transformer. 2. Find top-k clauses by cosine similarity to hypothesis. 3. Concatenate those clauses as the premise (truncated to 400 tokens). 4. Run NLI on the retrieved premise. This fixes the 512-token truncation bug entirely. """ st_model = models["chunker"].model # Reuse the sentence transformer # Embed everything in one batch all_texts = req.clauses + [req.hypothesis] embeddings = st_model.encode(all_texts, batch_size=32, show_progress_bar=False) clause_embeddings = embeddings[:len(req.clauses)] hypothesis_embedding = embeddings[-1].reshape(1, -1) # Cosine similarity similarities = cosine_similarity(hypothesis_embedding, clause_embeddings)[0] # Get top-k indices top_k = min(req.top_k, len(req.clauses)) top_indices = similarities.argsort()[-top_k:][::-1].tolist() top_clauses = [req.clauses[i] for i in top_indices] # Concatenate as premise, budget 400 tokens premise_parts = [] token_budget = 400 tokenizer = models["nli_tokenizer"] for clause in top_clauses: clause_tokens = len(tokenizer.encode(clause, add_special_tokens=False)) if clause_tokens <= token_budget: premise_parts.append(clause) token_budget -= clause_tokens else: # Truncate this clause to fit tokens = tokenizer.encode(clause, add_special_tokens=False)[:token_budget] premise_parts.append(tokenizer.decode(tokens)) break premise = " [...] ".join(premise_parts) # Run NLI inputs = models["nli_tokenizer"]( premise, req.hypothesis, return_tensors="pt", truncation=True, max_length=512 ) with torch.no_grad(): logits = models["nli_model"](**inputs).logits probs = torch.nn.functional.softmax(logits, dim=-1) class_id = torch.argmax(probs, dim=-1).item() return RetrievalNLIResult( label=models["nli_id2label"][class_id], score=round(probs[0][class_id].item(), 4), supporting_clauses=top_clauses, supporting_indices=top_indices ) # @app.post("/predict/nli_batch", response_model=BatchNLIResult) # def predict_nli_batch(req: BatchNLIRequest): # """Batch version: run N hypotheses against the same clause set.""" # results = [] # for hypothesis in req.hypotheses: # single_req = RetrievalNLIRequest( # clauses=req.clauses, # hypothesis=hypothesis, # top_k=req.top_k # ) # results.append(predict_nli_retrieval(single_req)) # return BatchNLIResult(results=results) # 2. Optimized NLI Batch (Huge Speedup) @app.post("/predict/nli_batch", response_model=BatchNLIResult) def predict_nli_batch(req: BatchNLIRequest): """Optimized: Embeds clauses once and reuses them for all hypotheses.""" if not req.clauses or not req.hypotheses: return BatchNLIResult(results=[]) st_model = models["chunker"].model tokenizer = models["nli_tokenizer"] nli_model = models["nli_model"] id2label = models["nli_id2label"] # Embed clauses once, hypotheses once with hf_lock: clause_embs = st_model.encode(req.clauses, batch_size=32, show_progress_bar=False) hyp_embs = st_model.encode(req.hypotheses, batch_size=32, show_progress_bar=False) # Vectorized similarity: (num_hyp x num_clauses) similarities = cosine_similarity(hyp_embs, clause_embs) results = [] for i, hyp in enumerate(req.hypotheses): sims = similarities[i] top_k = min(req.top_k, len(req.clauses)) top_indices = sims.argsort()[-top_k:][::-1].tolist() top_clauses = [req.clauses[j] for j in top_indices] # Premise construction (max 400 tokens) premise_parts, budget = [], 400 for c in top_clauses: with hf_lock: tokens = tokenizer.encode(c, add_special_tokens=False) if len(tokens) <= budget: premise_parts.append(c) budget -= len(tokens) else: with hf_lock: premise_parts.append(tokenizer.decode(tokens[:budget])) break # Inference with hf_lock: inputs = tokenizer(" [...] ".join(premise_parts), hyp, return_tensors="pt", truncation=True,max_length=512) with torch.no_grad(): logits = nli_model(**inputs).logits probs = torch.nn.functional.softmax(logits, dim=-1) class_id = torch.argmax(probs, dim=-1).item() results.append(RetrievalNLIResult( label=id2label[class_id], score=round(probs[0][class_id].item(), 4), supporting_clauses=top_clauses, supporting_indices=top_indices)) return BatchNLIResult(results=results) @app.post("/predict/ner_batch", response_model=BatchNERResult) def predict_ner_batch(req: BatchNERRequest): default_types = ["deposit percentage", "payment days", "cancellation fee percentage", "liability limit amount", "governing state", "notice period days", "confidentiality duration", "party name", "contract value","effective date"] labels = [l.lower() for l in (req.entity_types or default_types)] # GLiNER handles batches natively #batch_entities = models["ner"].predict_entities(req.texts, labels) batch_entities = [] for text in req.texts: # Force conversion to string and handle empty cases clean_text = str(text or "").strip() if not clean_text: batch_entities.append([]) continue try: with hf_lock: entities = models["ner"].predict_entities(clean_text, labels) batch_entities.append(entities) except Exception: batch_entities.append([]) # batch_entities = [ # models["ner"].predict_entities(text, labels) # for text in req.texts # ] return BatchNERResult(results=batch_entities) @app.post("/predict/minilm_embeddings",response_model=MiniLMEmbeddingResult) def get_minilm_embeddings(req: EmbeddingRequest): # Prevent huge requests from exhausting RAM if len(req.texts) > 1000: raise HTTPException( status_code=400, detail="Maximum 1000 texts per request" ) model = models["minilm_embeddings"] with hf_lock: embeddings = model.encode( req.texts, batch_size=32, show_progress_bar=False, convert_to_numpy=True, normalize_embeddings=True ) return MiniLMEmbeddingResult( embeddings=embeddings.tolist(), dimension=int(embeddings.shape[1]), model="sentence-transformers/all-MiniLM-L6-v2" )