Spaces:
Sleeping
Sleeping
import streamlit as st | |
from SearchEngine import searchengine | |
import logging | |
# Set up logging | |
logging.basicConfig(level=logging.INFO) | |
logger = logging.getLogger(__name__) | |
model_name = "all-MiniLM-L6-v2" | |
collection_name = "docs" | |
search_engine = searchengine(model_name, collection_name) | |
def main(): | |
st.title("ChromaDB Search Engine") | |
st.sidebar.header("Add Document") | |
text_input = st.sidebar.text_area("Text") | |
metadata_input = {'type':st.sidebar.text_input("Type")} | |
add_button = st.sidebar.button("Save") | |
if add_button: | |
document_id = str(search_engine.count() + 1) | |
search_engine.add(text_input, metadata_input, document_id) | |
st.sidebar.success(f"Document added with ID: {document_id}") | |
st.sidebar.header("Search") | |
query = st.sidebar.text_input("Query") | |
search_button = st.sidebar.button("Search") | |
if search_button: | |
results = search_engine.query(query) | |
st.subheader("Search Results:") | |
logging.info("result :") | |
logging.info(results) | |
ids = results['ids'][0] | |
distances = results['distances'][0] | |
metadatas = results['metadatas'][0] | |
documents = results['documents'][0] | |
for index, id in enumerate(ids): | |
st.write(f"Document ID: {id}, Metadata: {metadatas[index]}") | |
st.write(f"Text: {documents[index]}") | |
st.write(f"distance: {distances[index]}") | |
st.markdown("---") | |
if __name__ == "__main__": | |
main() | |