Sebastian Kułaga commited on
Commit
dec29b0
1 Parent(s): cc07cb1
Files changed (2) hide show
  1. app.py +2 -2
  2. routers/distilbert.py +2 -2
app.py CHANGED
@@ -1,11 +1,11 @@
1
  from fastapi import FastAPI
2
 
3
- from routers.distilbert import predict
4
 
5
  # Create a new FastAPI app instance
6
  app = FastAPI()
7
 
8
- app.include_router(predict)
9
 
10
  @app.get("/")
11
  def greet_json():
 
1
  from fastapi import FastAPI
2
 
3
+ from routers.distilbert import prediction
4
 
5
  # Create a new FastAPI app instance
6
  app = FastAPI()
7
 
8
+ app.include_router(prediction)
9
 
10
  @app.get("/")
11
  def greet_json():
routers/distilbert.py CHANGED
@@ -3,12 +3,12 @@ from src.distilbert import Distilbert
3
 
4
 
5
  MODEL_PATH = "SebaK13/DistilBERT-finetuned-customer-queries-balanced"
6
- predict = APIRouter(prefix="/predict", tags=["predict"])
7
  distilbert_service = Distilbert(model_path=MODEL_PATH)
8
 
9
 
10
 
11
- @predict.get("/predict")
12
  def predict(query: str):
13
  output = distilbert_service.predict_query_type(query)
14
  return {"label": output}
 
3
 
4
 
5
  MODEL_PATH = "SebaK13/DistilBERT-finetuned-customer-queries-balanced"
6
+ prediction = APIRouter(prefix="/predict", tags=["predict"])
7
  distilbert_service = Distilbert(model_path=MODEL_PATH)
8
 
9
 
10
 
11
+ @prediction.get("/predict")
12
  def predict(query: str):
13
  output = distilbert_service.predict_query_type(query)
14
  return {"label": output}