Spaces:
Sleeping
Sleeping
File size: 1,508 Bytes
0737a07 9d5b5fa 8699e3e 8dd99a4 8699e3e 8dd99a4 0737a07 6457659 f543029 8dd99a4 f543029 adbfa49 f543029 c7ac645 f543029 adbfa49 f543029 8fe1189 6457659 6b54693 8fe1189 6b54693 f543029 acc18ac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
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()
|