GPT-Docker / app /utils.py
heikowagner's picture
move elements
98639ab
raw
history blame
1.34 kB
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