File size: 1,393 Bytes
0ca7ed3 d27d5a2 0ca7ed3 d36aadb 0ca7ed3 |
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 |
import os
import streamlit as st
import requests
from typing import List
import json
import socket
from urllib3.connection import HTTPConnection
from app import embed_documents, retrieve_documents
# API_BASE_URL = os.environ.get("API_BASE_URL")
embeddings_model_name = "all-MiniLM-L6-v2"
persist_directory = "db"
model = "tiiuae/falcon-7b-instruct"
from constants import CHROMA_SETTINGS
import chromadb
def list_of_collections():
client = chromadb.Client(CHROMA_SETTINGS)
return (client.list_collections())
def get_collection_names():
collections = list_of_collections()
return [collection.name for collection in collections]
if __name__ == '__main__':
st.title("PrivateGPT App: Document Embedding and Retrieval")
# Document upload section
st.header("Document Upload")
files = st.file_uploader("Upload document", accept_multiple_files=True)
# collection_name = st.text_input("Collection Name") not working for some reason
if st.button("Embed"):
embed_documents(files, "collection_name")
# Query section
st.header("Document Retrieval")
collection_names = get_collection_names()
selected_collection = st.selectbox("Select a document", collection_names)
if selected_collection:
query = st.text_input("Query")
if st.button("Retrieve"):
retrieve_documents(query, selected_collection)
|