Spaces:
Sleeping
Sleeping
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) |