Update app.py

app.py

@@ -3,6 +3,8 @@ from pydantic import BaseModel
 from model.model import predict_pipeline
 from model.model import __version__ as model_version
 
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import TextClassificationPipeline
 
 app = FastAPI()
 
@@ -14,6 +16,9 @@ class TextIn(BaseModel):
 class PredictionOut(BaseModel):
     language: str
 
+class TopicClassificationOut(BaseModel):
+    result: str
+
 
 @app.get("/")
 def home():
@@ -23,4 +28,15 @@ def home():
 @app.post("/predict", response_model=PredictionOut)
 def predict(payload: TextIn):
     language = predict_pipeline(payload.text)
-    return {"language": language}
+    return {"language": language}
+
+@app.post("/TopicClassification", response_model=TopicClassificationOut)
+def TopicClassification(payload: TextIn):
+    model_name = 'lincoln/flaubert-mlsum-topic-classification'
+    
+    loaded_tokenizer = AutoTokenizer.from_pretrained(model_name)
+    loaded_model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    
+    nlp = TextClassificationPipeline(model=loaded_model, tokenizer=loaded_tokenizer)
+    result = nlp(payload.text, truncation=True)
+    return {"result": result}