Spaces:

HinaCortus
/

API_KnowYourAudience

Runtime error

Update app.py

d8f90a3 about 1 year ago

1.36 kB

	from fastapi import FastAPI
	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()


	class TextIn(BaseModel):
	text: str


	class PredictionOut(BaseModel):
	language: str

	class TopicClassificationOut(BaseModel):
	result: str


	@app.get("/")
	def home():
	return {"health_check": "OK", "model_version": model_version}


	@app.post("/predict", response_model=PredictionOut)
	def predict(payload: TextIn):
	language = predict_pipeline(payload.text)
	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)

	print(payload.text)
	text = re.sub(r'[!@#$(),\n"%^*?\:;~`0-9]', " ", payload.text)
	text = re.sub(r"[[]]", " ", text)
	text = text.lower()

	result = nlp(text, truncation=True)
	return {"result": result}