chatroles-vdb / app.py
pond918
feat: embed array of text
130076a
import gradio as gr
import json
import faiss_vdb as vdb
def upsert_role(roleId, data):
"""
body: { meta, content }
"""
text = json.dumps(data['content'])
meta = data['meta']
meta['id'] = roleId
cnt = vdb.upsert(text, meta)
return cnt
i_upsert = gr.Interface(fn=upsert_role, inputs=["text", "json"],
outputs="text")
# gradio_interface.launch()
def search_roles(content, size):
if size is None: size = 4
size = int(size)
text = json.dumps(content)
docs = vdb.search(text, size)
return docs
i_search = gr.Interface(
fn=search_roles,
inputs=["json", "number"],
outputs="json"
)
def embed_texts(texts):
arr = [vdb.embed_text(t)[0].tolist() for t in texts]
json_arr = json.dumps(arr)
return json_arr
i_embed = gr.Interface(
fn=embed_texts,
inputs='json',
outputs="json"
)
demo = gr.TabbedInterface([i_upsert, i_search, i_embed], [
"upsert role", "search roles", "embed text"])
if __name__ == "__main__":
demo.launch()