File size: 770 Bytes
af9d72e
 
 
 
 
54bafc8
af9d72e
 
 
 
 
 
 
 
 
54bafc8
af9d72e
 
 
54bafc8
 
 
 
 
 
 
af9d72e
 
 
 
 
 
 
 
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
from fastapi import FastAPI
from pydantic import BaseModel
from model.model import LLM
import torch

app = FastAPI(docs_url="/swagger-ui.html")

class InputText(BaseModel):
    text: str

# "bigscience/bloomz-1b1"
model_tag = "facebook/opt-125m"
model = LLM(model_name = model_tag,
            device = "cuda" if torch.cuda.is_available() else "cpu")

"""
@app.get("/")
async def root():
    return {"message": "Technical challenge OK"}
"""

@app.get("/")
async def docs_redirect():
    response = RedirectResponse(url='/swagger-ui.html')
    return response


@app.post("/language-detection")
def language_detection(text):
    return {"language": model.language_detection(text)}

@app.post("/entity-recognition")
def ner(text):
    return model.entity_recognition(text)