Spaces:
Running
Running
import streamlit as st | |
import chromadb | |
from chromadb.utils import embedding_functions | |
from sentence_transformers import SentenceTransformer | |
client = chromadb.PersistentClient(path="./chromadb/") | |
MODEL_NAME: str = "mixedbread-ai/mxbai-embed-large-v1" # ~ 0.5 gb | |
COLLECTION_NAME: str = "scheme" | |
EMBEDDING_FUNC = embedding_functions.SentenceTransformerEmbeddingFunction(model_name=MODEL_NAME) | |
schemer = client.get_collection( | |
name=COLLECTION_NAME, | |
embedding_function=EMBEDDING_FUNC, | |
) | |
DATA_AVAL: bool = schemer.count() > 0 | |
APP_NAME: str = "Groove-GPT" | |
st.title(APP_NAME) | |
st.write("Data Avaliable: ", DATA_AVAL) | |
user_question: str = st.text_area("Enter your groovy questions here") | |
access_key: str = st.text_input("Enter your gpt key here", type="password") | |
if st.button('Query Database'): | |
if (user_question == ""): | |
st.stop() | |
else: | |
st.header("Results") | |
# Perform the Chromadb query. | |
results = schemer.query( | |
query_texts=[user_question], | |
n_results=10, | |
include = ['documents'] | |
) | |
st.write(results) |