app.py

import gradio as gr
from gradio.mix import Parallel
from transformers import pipeline
from sentence_transformers import SentenceTransformer, util
from rembg import remove


generator = pipeline('text-generation', model='gpt2')
model_txt_embedding = SentenceTransformer("sentence-transformers/all-mpnet-base-v2")


# 文本向量
def embedding(txt):
    return str(model_txt_embedding.encode(txt, convert_to_tensor=True, convert_to_numpy=True).tolist())

# 文本相似度
def embedding_sim(text1, text2):
    embeddings1 = model_txt_embedding.encode(text1, convert_to_tensor=True)
    embeddings2 = model_txt_embedding.encode(text2, convert_to_tensor=True)
    return str(util.cos_sim(embeddings1, embeddings2).item())

def remove_bg(image):
    # return image
    return remove(image)

tab1 = gr.Interface(fn=embedding, inputs="text", outputs="text")
tab2 = gr.Interface(fn=embedding_sim, inputs=["text", "text"], outputs="text")
tab3 = gr.Interface(fn=remove_bg, inputs="image", outputs="image")

app = gr.TabbedInterface([tab1, tab2, tab3], ["文本向量", "文本相似度", "抠图"])
app.queue()
app.launch()