|
import torch |
|
from fastapi import FastAPI, Depends, status |
|
from fastapi.responses import PlainTextResponse |
|
from pydantic import BaseModel |
|
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM |
|
|
|
import time |
|
from typing import Dict, List, Optional |
|
|
|
import jwt |
|
from decouple import config |
|
from fastapi import Request, HTTPException |
|
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials |
|
|
|
JWT_SECRET = config("secret") |
|
JWT_ALGORITHM = config("algorithm") |
|
|
|
app = FastAPI() |
|
app.ready = False |
|
|
|
device = ("cuda" if torch.cuda.is_available() else "cpu") |
|
tokenizer = AutoTokenizer.from_pretrained('vblagoje/bart_lfqa') |
|
model = AutoModelForSeq2SeqLM.from_pretrained('vblagoje/bart_lfqa').to(device) |
|
_ = model.eval() |
|
|
|
|
|
class JWTBearer(HTTPBearer): |
|
def __init__(self, auto_error: bool = True): |
|
super(JWTBearer, self).__init__(auto_error=auto_error) |
|
|
|
async def __call__(self, request: Request): |
|
credentials: HTTPAuthorizationCredentials = await super(JWTBearer, self).__call__(request) |
|
if credentials: |
|
if not credentials.scheme == "Bearer": |
|
raise HTTPException(status_code=403, detail="Invalid authentication scheme.") |
|
if not self.verify_jwt(credentials.credentials): |
|
raise HTTPException(status_code=403, detail="Invalid token or expired token.") |
|
return credentials.credentials |
|
else: |
|
raise HTTPException(status_code=403, detail="Invalid authorization code.") |
|
|
|
def verify_jwt(self, jwtoken: str) -> bool: |
|
isTokenValid: bool = False |
|
|
|
try: |
|
payload = decodeJWT(jwtoken) |
|
except: |
|
payload = None |
|
if payload: |
|
isTokenValid = True |
|
return isTokenValid |
|
|
|
|
|
def token_response(token: str): |
|
return { |
|
"access_token": token |
|
} |
|
|
|
|
|
def signJWT(user_id: str) -> Dict[str, str]: |
|
payload = { |
|
"user_id": user_id, |
|
"expires": time.time() + 6000 |
|
} |
|
token = jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM) |
|
|
|
return token_response(token) |
|
|
|
|
|
def decodeJWT(token: str) -> dict: |
|
try: |
|
decoded_token = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM]) |
|
return decoded_token if decoded_token["expires"] >= time.time() else None |
|
except: |
|
return {} |
|
|
|
|
|
class LFQAParameters(BaseModel): |
|
min_length: int = 50 |
|
max_length: int = 250 |
|
do_sample: bool = False |
|
early_stopping: bool = True |
|
num_beams: int = 8 |
|
temperature: float = 1.0 |
|
top_k: float = None |
|
top_p: float = None |
|
no_repeat_ngram_size: int = 3 |
|
num_return_sequences: int = 1 |
|
|
|
|
|
class InferencePayload(BaseModel): |
|
model_input: str |
|
parameters: Optional[LFQAParameters] = LFQAParameters() |
|
|
|
|
|
@app.on_event("startup") |
|
def startup(): |
|
app.ready = True |
|
|
|
|
|
@app.get("/healthz") |
|
def healthz(): |
|
if app.ready: |
|
return PlainTextResponse("ok") |
|
return PlainTextResponse("service unavailable", status_code=status.HTTP_503_SERVICE_UNAVAILABLE) |
|
|
|
|
|
@app.post("/generate/", dependencies=[Depends(JWTBearer())]) |
|
def generate(context: InferencePayload): |
|
|
|
model_input = tokenizer(context.model_input, truncation=True, padding=True, return_tensors="pt") |
|
param = context.parameters |
|
generated_answers_encoded = model.generate(input_ids=model_input["input_ids"].to(device), |
|
attention_mask=model_input["attention_mask"].to(device), |
|
min_length=param.min_length, |
|
max_length=param.max_length, |
|
do_sample=param.do_sample, |
|
early_stopping=param.early_stopping, |
|
num_beams=param.num_beams, |
|
temperature=param.temperature, |
|
top_k=param.top_k, |
|
top_p=param.top_p, |
|
no_repeat_ngram_size=param.no_repeat_ngram_size, |
|
num_return_sequences=param.num_return_sequences) |
|
answers = tokenizer.batch_decode(generated_answers_encoded, skip_special_tokens=True, |
|
clean_up_tokenization_spaces=True) |
|
results = [] |
|
for answer in answers: |
|
results.append({"generated_text": answer}) |
|
return results |
|
|
