app.py

import torch
import numpy as np
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

models = {}

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)")

    # 4. Legal BERT embeddings
    print("Loading legal BERT embeddings model...")
    emb_model, emb_tokenizer = load_quantized_bert("nlpaueb/bert-base-uncased-contracts", num_labels=None)
    models["emb_tokenizer"] = emb_tokenizer
    models["emb_model"] = emb_model
    #models["emb_model"].eval()
    print("✓ legal BERT loaded(int8)")

    # 4b. Text explanation / summarization model (Flan‑T5 small, Float16)
    print("Loading text explanation/summarization model...")
    explain_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
    explain_model = AutoModelForSeq2SeqLM.from_pretrained(
        "google/flan-t5-small",
        torch_dtype=torch.float16       # half‑precision, good tradeoff
    ).eval()
    models["explain_model"] = explain_model
    models["explain_tokenizer"] = explain_tokenizer
    print("✓ explain/summarize model loaded (flan-t5-small, float16)")

    # 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)
    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 EmbeddingRequest(BaseModel):
    texts: List[str]

class ChunkRequest(BaseModel):
    text: str
    percentile_threshold: float = 75.0
    context_window: int = 1
    min_chunk_size: int = 2


class ExplanationRequest(BaseModel):
    text: str
    mode: str = "explain"   # "summarize" or "explain"


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]




# ---------- 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)
    )


def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output.last_hidden_state
    mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * mask_expanded, 1) / torch.clamp(mask_expanded.sum(1), min=1e-9)


@app.post("/predict/embeddings", response_model=EmbeddingResult)
def get_embeddings(req: EmbeddingRequest):
    encoded = models["emb_tokenizer"](
        req.texts, padding=True, truncation=True, return_tensors="pt"
    )
    with torch.no_grad():
        outputs = models["emb_model"](**encoded)
    embeddings = mean_pooling(outputs, encoded["attention_mask"])
    return EmbeddingResult(embeddings=embeddings.tolist())


@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/explain")
def explain_text(req: ExplanationRequest):
    tokenizer = models["explain_tokenizer"]
    model = models["explain_model"]

    # FLAN-T5 models fine-tuned on summarization require the "summarize: " prefix
    input_text = f"summarize: {req.text}"
    
    # If the user asks for an 'explain', we can still frame it as an intensive summary
    if req.mode == "explain":
        input_text = f"summarize in detail: {req.text}"

    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=512)
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            num_beams=5,
            length_penalty=2.0,      # Encourage longer generation
            no_repeat_ngram_size=3,  # Prevent repetition
            early_stopping=True
        )
    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return {"mode": req.mode, "generated_text": result}




@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)