Spaces:
Running
Running
# main.py | |
import os | |
import tempfile | |
import streamlit as st | |
from question import chat_with_doc | |
from langchain.embeddings import HuggingFaceInferenceAPIEmbeddings | |
from langchain.vectorstores import SupabaseVectorStore | |
from supabase import Client, create_client | |
from stats import get_usage | |
supabase_url = st.secrets.SUPABASE_URL | |
supabase_key = st.secrets.SUPABASE_KEY | |
openai_api_key = st.secrets.openai_api_key | |
anthropic_api_key = st.secrets.anthropic_api_key | |
hf_api_key = st.secrets.hf_api_key | |
supabase: Client = create_client(supabase_url, supabase_key) | |
self_hosted = st.secrets.self_hosted | |
username = st.secrets.username | |
# embeddings = OpenAIEmbeddings(openai_api_key=openai_api_key) | |
embeddings = HuggingFaceInferenceAPIEmbeddings( | |
api_key=hf_api_key, | |
model_name="BAAI/bge-large-en-v1.5" | |
) | |
vector_store = SupabaseVectorStore(supabase, embeddings, query_name='match_documents', table_name="documents") | |
models = ["meta-llama/Llama-2-70b-chat-hf", "mistralai/Mixtral-8x7B-Instruct-v0.1"] | |
if openai_api_key: | |
models += ["gpt-3.5-turbo", "gpt-4"] | |
if anthropic_api_key: | |
models += ["claude-v1", "claude-v1.3", | |
"claude-instant-v1-100k", "claude-instant-v1.1-100k"] | |
# Set the theme | |
st.set_page_config( | |
page_title="Securade.ai - Safety Copilot", | |
page_icon="https://securade.ai/favicon.ico", | |
layout="centered", | |
initial_sidebar_state="collapsed", | |
menu_items={ | |
"About": "# Securade.ai Safety Copilot v0.1\n [https://securade.ai](https://securade.ai)", | |
"Get Help" : "https://securade.ai", | |
"Report a Bug": "mailto:[email protected]" | |
} | |
) | |
st.title("👷♂️ Safety Copilot 🦺") | |
st.markdown("Chat with your personal safety assistant about any health & safety related queries.") | |
st.markdown("Up-to-date with latest OSH regulations for Singapore, Indonesia, Malaysia & other parts of Asia.") | |
st.markdown("---\n\n") | |
# Initialize session state variables | |
if 'model' not in st.session_state: | |
st.session_state['model'] = "meta-llama/Llama-2-70b-chat-hf" | |
if 'temperature' not in st.session_state: | |
st.session_state['temperature'] = 0.1 | |
if 'chunk_size' not in st.session_state: | |
st.session_state['chunk_size'] = 500 | |
if 'chunk_overlap' not in st.session_state: | |
st.session_state['chunk_overlap'] = 0 | |
if 'max_tokens' not in st.session_state: | |
st.session_state['max_tokens'] = 500 | |
if 'username' not in st.session_state: | |
st.session_state['username'] = username | |
stats = str(get_usage(supabase)) | |
chat_with_doc(st.session_state['model'], vector_store, stats_db=supabase, stats=stats) | |
st.markdown("---\n\n") |