from fastapi import FastAPI
from pydantic import BaseModel
import onnxruntime as ort
import onnxruntime
print(f"ORT version: {onnxruntime.__version__}")
import numpy as np
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download
import os
import shutil

app = FastAPI()
MODEL_REPO = "Sandeep120205/agent-shield-mdeberta"
THRESHOLD = 0.85
MAX_LEN = 128

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained("microsoft/mdeberta-v3-base")

print("Loading model...")
cache_dir = os.path.expanduser("~/.cache/huggingface/hub")
if os.path.exists(cache_dir):
    shutil.rmtree(cache_dir)
    print("Cache cleared.")

model_path = hf_hub_download(repo_id=MODEL_REPO, filename="model_fp32.onnx", force_download=True)
session = ort.InferenceSession(model_path)
print(f"ONNX path: {model_path}")
print(f"ONNX size: {os.path.getsize(model_path)}")
expected_inputs = [inp.name for inp in session.get_inputs()]
print(f"Model loaded. Inputs: {expected_inputs}")

class PredictRequest(BaseModel):
    prompt: str

@app.get("/health")
def health():
    return {"status": "ok"}

@app.post("/predict")
def predict(req: PredictRequest):
    inputs = tokenizer(
        req.prompt,
        return_tensors="np",
        truncation=True,
        max_length=MAX_LEN,
        padding="max_length"
    )
    filtered = {k: v for k, v in inputs.items() if k in expected_inputs}
    outputs = session.run(None, filtered)
    logits = outputs[0][0]
    exp = np.exp(logits - np.max(logits))
    probs = exp / exp.sum()
    confidence = float(probs[1])
    is_injection = confidence > THRESHOLD
    return {"is_injection": is_injection, "confidence": confidence}