Spaces:
Build error
Build error
| import streamlit as st | |
| import latex2markdown | |
| from langchain.docstore.document import Document | |
| import chromadb | |
| from chromadb.config import Settings | |
| import load_model | |
| from load_model import load_embedding | |
| from load_vectors import load_from_file, load_and_split, create_and_add, load_from_web | |
| persist_directory = load_model.persist_directory | |
| def format_document(document: Document): | |
| """TODO: Implement a nice style""" | |
| return document.dict() | |
| def format_result_set(result): | |
| st.write(latex2markdown.LaTeX2Markdown(result["result"]).to_markdown()) | |
| agree = st.checkbox('Show source documents') | |
| source_documents = result["source_documents"] | |
| if agree: | |
| st.write('Source Documents:') | |
| for document in source_documents: | |
| st.write(format_document(document)) | |
| #@st.cache_resource | |
| def get_chroma_client(): | |
| return chromadb.Client(Settings(chroma_db_impl="duckdb+parquet", | |
| persist_directory=persist_directory | |
| )) | |
| #@st.cache_data | |
| def retrieve_collections(): | |
| client = get_chroma_client() | |
| all_collections = client.list_collections() | |
| print(all_collections) | |
| print(all_collections[0].metadata) | |
| collections = tuple( [{'name': collection.name, 'model_name': collection.metadata['model_name']} for collection in all_collections] ) | |
| return collections | |
| def load_files(): | |
| client = get_chroma_client() | |
| option = st.radio( | |
| "", | |
| options=["Add Documents", "Start new collection"], | |
| ) | |
| collections = retrieve_collections() | |
| if option == "Add Documents": | |
| selected_collection = st.selectbox( | |
| 'Add to exsisting collection or create a new one', | |
| collections ) | |
| if st.button('Delete Collection (⚠️ This is destructive and not reversible)'): | |
| client.delete_collection(name=selected_collection["name"]) | |
| #retrieve_collections.clear() | |
| collections = retrieve_collections() | |
| option = st.radio( | |
| "", | |
| options=["Upload Files from Local", "Upload Files from Web"], | |
| ) | |
| if option == "Upload Files from Local": | |
| st.write('Source Documents:') | |
| uploaded_files = st.file_uploader("Choose a PDF file", accept_multiple_files=True) | |
| chunk_size = st.text_area('chunk Size:', 1000) | |
| if st.button('Upload'): | |
| docs = load_from_file(uploaded_files) | |
| sub_docs = load_and_split(docs, chunk_size=int(chunk_size)) | |
| create_and_add(selected_collection["name"], sub_docs, selected_collection['model_name'], {"model_name": selected_collection['model_name']}) | |
| uploaded_files=None | |
| st.write("Upload succesful") | |
| else: | |
| st.write('Urls of Source Documents (Comma separated):') | |
| urls = chunk_size = st.text_area('Urls:', '') | |
| chunk_size = st.text_area('chunk Size:', 1000) | |
| urls = urls.replace(",", "" ).replace('"', "" ).split(',') | |
| if st.button('Upload'): | |
| docs = load_from_web(urls) | |
| sub_docs = load_and_split(docs, chunk_size=int(chunk_size)) | |
| print(selected_collection['model_name']) | |
| create_and_add(selected_collection["name"], sub_docs, selected_collection['model_name'], {"model_name": selected_collection['model_name']}) | |
| uploaded_files=None | |
| st.write("Upload succesful") | |
| else: | |
| collection = st.text_area('Name of your new collection:', '') | |
| model_name = st.text_area('Choose the embedding function:', "hkunlp/instructor-large") | |
| if st.button('Create'): | |
| if len(collection)>3: | |
| ef = load_embedding(model_name) | |
| metadata= {"loaded_docs":[], "Subject":"Terms Example", "model_name": ef.model_name} | |
| client.create_collection(collection, embedding_function=ef, metadata=metadata) | |
| # retrieve_collections.clear() | |
| st.write("Collection " +collection+" succesfully created.") |