File size: 1,337 Bytes
7009660
 
 
fbb697c
 
 
1f84a9a
4f0dc21
fbb697c
7009660
 
 
 
 
 
 
 
 
 
 
 
 
 
 
73d6caa
fbb697c
 
 
 
83629b5
fbb697c
 
1f84a9a
aeb550e
98639ab
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
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()
    collections = tuple( [{'name': collection.name, 'model_name': collection.metadata['model_name'], "metadata": collection.metadata} for collection in all_collections] )
    return collections