pipeline / main.py
lalalic's picture
Update main.py
52909e5 verified
import os, json, logging
import requests
from transformers import pipeline
from flask import Flask, request, jsonify
app=Flask(__name__)
@app.route("/")
def hello():
return "welcome!"
def parse_params():
input=request.args.get('input')
kargs=request.args.get('kargs')
try:
input = json.load(input)
except:
pass
try:
kargs = json.load(kargs)
except:
kargs={}
return input, kargs
from huggingface_hub import HfApi
api=HfApi()
@app.route("/search")
def search():
args=request.args.to_dict()
models = api.list_models(**args)
return jsonify(models)
@app.route("/task_list")
def tasks():
return [item.strip() for item in
'''audio-classification
automatic-speech-recognition
conversational
depth-estimation
document-question-answering
feature-extraction
fill-mask
image-classification
image-feature-extraction
image-segmentation
image-to-image
image-to-text
mask-generation
object-detection
question-answering
summarization
table-question-answering
text2text-generation
text-classification (alias sentiment-analysis available)
text-generation
text-to-audio (alias text-to-speech available)
token-classification (alias ner available)
translation
translation_xx_to_yy
video-classification
visual-question-answering
zero-shot-classification
zero-shot-image-classification
zero-shot-audio-classification
zero-shot-object-detection'''.split("\n")
]
@app.route("/<task>")
def run_task(task):
(input, kargs)=parse_params()
pipe=pipeline(task, **kargs)
return pipe(input)
@app.route("/<user>/<model>")
def run_model():
(input, kargs)=parse_params()
pipe=pipeline(model=f'{user}/{model}', **kargs)
return pipe(input)
@app.route("/<task>/<user>/<model>")
def run_task_model():
(input, kargs)=parse_params()
pipe=pipeline(task, model=f'{user}/{model}', **kargs)
return pipe(input)
logging.info("xtts is ready")