HinaCortus commited on
Commit
89dd83c
1 Parent(s): 4bd7b39

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +17 -1
app.py CHANGED
@@ -3,6 +3,8 @@ from pydantic import BaseModel
3
  from model.model import predict_pipeline
4
  from model.model import __version__ as model_version
5
 
 
 
6
 
7
  app = FastAPI()
8
 
@@ -14,6 +16,9 @@ class TextIn(BaseModel):
14
  class PredictionOut(BaseModel):
15
  language: str
16
 
 
 
 
17
 
18
  @app.get("/")
19
  def home():
@@ -23,4 +28,15 @@ def home():
23
  @app.post("/predict", response_model=PredictionOut)
24
  def predict(payload: TextIn):
25
  language = predict_pipeline(payload.text)
26
- return {"language": language}
 
 
 
 
 
 
 
 
 
 
 
 
3
  from model.model import predict_pipeline
4
  from model.model import __version__ as model_version
5
 
6
+ from transformers import AutoTokenizer, AutoModelForSequenceClassification
7
+ from transformers import TextClassificationPipeline
8
 
9
  app = FastAPI()
10
 
 
16
  class PredictionOut(BaseModel):
17
  language: str
18
 
19
+ class TopicClassificationOut(BaseModel):
20
+ result: str
21
+
22
 
23
  @app.get("/")
24
  def home():
 
28
  @app.post("/predict", response_model=PredictionOut)
29
  def predict(payload: TextIn):
30
  language = predict_pipeline(payload.text)
31
+ return {"language": language}
32
+
33
+ @app.post("/TopicClassification", response_model=TopicClassificationOut)
34
+ def TopicClassification(payload: TextIn):
35
+ model_name = 'lincoln/flaubert-mlsum-topic-classification'
36
+
37
+ loaded_tokenizer = AutoTokenizer.from_pretrained(model_name)
38
+ loaded_model = AutoModelForSequenceClassification.from_pretrained(model_name)
39
+
40
+ nlp = TextClassificationPipeline(model=loaded_model, tokenizer=loaded_tokenizer)
41
+ result = nlp(payload.text, truncation=True)
42
+ return {"result": result}