Spaces:
Running
Running
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() |