Spaces:
Sleeping
Sleeping
File size: 2,718 Bytes
ab05059 9f1fab5 ab05059 9f1fab5 ab05059 9f1fab5 ab05059 9f1fab5 ab05059 9f1fab5 ab05059 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 |
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import pipeline
import uvicorn
app = FastAPI(title="PrivAware API", version="2.0")
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
print("Loading phishing model...")
phishing_classifier = pipeline(
"text-classification",
model="V3d4nt7/privaware-phishing-detector"
)
print("Loading policy model...")
policy_classifier = pipeline(
"text-classification",
model="V3d4nt7/privaware-policy-classifier"
)
print("Both models loaded!")
class URLRequest(BaseModel):
url: str
class PolicyRequest(BaseModel):
text: str
@app.get("/")
def root():
return {"status": "ok", "message": "PrivAware API v2 running — both models live"}
@app.get("/health")
def health():
return {"status": "healthy", "models": ["phishing-detector", "policy-classifier"]}
@app.post("/scan-url")
def scan_url(request: URLRequest):
if not request.url:
raise HTTPException(status_code=400, detail="URL is required")
try:
result = phishing_classifier(request.url[:512])[0]
label = result["label"]
confidence = round(result["score"] * 100, 2)
risk_score = round(confidence) if label == "PHISHING" else round(100 - confidence)
return {
"url": request.url,
"label": label,
"confidence": confidence,
"risk_score": risk_score,
"is_phishing": label == "PHISHING"
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/scan-policy")
def scan_policy(request: PolicyRequest):
if not request.text:
raise HTTPException(status_code=400, detail="Text is required")
try:
result = policy_classifier(request.text[:512])[0]
label = result["label"]
confidence = round(result["score"] * 100, 2)
risk_map = {"SAFE": 10, "RISKY": 55, "DECEPTIVE": 90}
risk_score = risk_map.get(label, 50)
return {
"label": label,
"confidence": confidence,
"risk_score": risk_score,
"is_deceptive": label == "DECEPTIVE"
}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/scan-full")
def scan_full(request: URLRequest):
url_result = scan_url(request)
combined_score = url_result["risk_score"]
return {
"url": request.url,
"phishing": url_result,
"combined_risk_score": combined_score
}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)
|