Spaces:
Runtime error
Runtime error
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 | |
def home(): | |
return {"health_check": "OK", "model_version": model_version} | |
def predict(payload: TextIn): | |
language = predict_pipeline(payload.text) | |
return {"language": language} | |
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 = "Le Bayern Munich prend la grenadine." | |
text = re.sub(r"[[]]", " ", text) | |
text = text.lower() | |
result = nlp(text, truncation=True) | |
return {"result": result} | |