Spaces:
Sleeping
Sleeping
File size: 3,715 Bytes
d2237d8 6a362a8 d2237d8 e1ec338 d2237d8 e1ec338 d2237d8 75c8918 d2237d8 75c8918 5c6714e 41b77f5 8d1af45 41b77f5 8d1af45 41b77f5 ef88834 8d1af45 75c8918 8d1af45 75c8918 d2237d8 c43f30f 6a362a8 c43f30f 6a362a8 c43f30f |
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 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 |
from fastapi import FastAPI, Depends, HTTPException, Request
from fastapi.responses import HTMLResponse # ✅ Add this line
from fastapi.security import APIKeyHeader
from pydantic import BaseModel
from model import ner_pipeline
import logging
import time
import json
import os
import secrets
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
app = FastAPI()
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
#API_KEY_FILE = "api_keys.json"
API_KEY_FILE = "/tmp/api_keys.json"
token_header = APIKeyHeader(name="X-API-KEY")
if os.path.exists(API_KEY_FILE):
with open(API_KEY_FILE, "r") as f:
API_KEYS_STORE = json.load(f)
else:
API_KEYS_STORE = {"users": {}}
with open(API_KEY_FILE, "w") as f:
json.dump(API_KEYS_STORE, f)
API_KEYS = API_KEYS_STORE.get("users", {})
ADMIN_KEY = os.getenv("ADMIN_KEY")
class TextRequest(BaseModel):
text: str
class RegisterRequest(BaseModel):
label: str = "user"
def save_keys():
API_KEYS_STORE["users"] = API_KEYS
with open(API_KEY_FILE, "w") as f:
json.dump(API_KEYS_STORE, f, indent=2)
def verify_token(x_api_key: str = Depends(token_header)):
if x_api_key not in API_KEYS:
raise HTTPException(status_code=403, detail="Unauthorized")
return x_api_key
"""
@app.post("/ner")
@limiter.limit("10/minute")
def ner_predict(request: TextRequest, api_key: str = Depends(verify_token), req: Request = None):
logger.info("Received NER request from IP: %s", get_remote_address(req))
predictions = ner_pipeline(request.text)
API_KEYS[api_key]["usage_count"] = API_KEYS[api_key].get("usage_count", 0) + 1
save_keys()
return {
"entities": predictions,
"usage": API_KEYS[api_key]["usage_count"]
}
"""
from fastapi import Request
@app.post("/ner")
@limiter.limit("10/minute")
def ner_predict(body: TextRequest, request: Request, api_key: str = Depends(verify_token)):
logger.info("NER request from IP: %s", get_remote_address(request))
raw_predictions = ner_pipeline(body.text)
# Use original character spans to extract clean entity words
predictions = [
{
**ent,
"score": float(ent["score"]),
"word": body.text[ent["start"]:ent["end"]].strip() # ✅ Clean spaces
}
for ent in raw_predictions
]
API_KEYS[api_key]["usage_count"] = API_KEYS[api_key].get("usage_count", 0) + 1
save_keys()
return {
"entities": predictions,
"usage": API_KEYS[api_key]["usage_count"]
}
@app.post("/register_user")
def register_user(request: RegisterRequest, x_api_key: str = Depends(token_header)):
if x_api_key != ADMIN_KEY:
raise HTTPException(status_code=403, detail="Admin access required")
new_key = secrets.token_urlsafe(32)
API_KEYS[new_key] = {"usage_count": 0, "label": request.label}
save_keys()
return {"message": "User registered", "api_key": new_key}
@app.get("/list_users")
def list_users(x_api_key: str = Depends(token_header)):
if x_api_key != ADMIN_KEY:
raise HTTPException(status_code=403, detail="Admin access required")
return {"users": API_KEYS}
@app.get("/", include_in_schema=False)
def root():
return HTMLResponse("""
<html>
<head><title>NER API</title></head>
<body>
<h1>✅ NER API is running!</h1>
<p>Visit <a href='/docs'>/docs</a> to try the API.</p>
</body>
</html>
""")
|