add logo

app.py

@@ -1,16 +1,19 @@
+import base64
+from pathlib import Path
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from langchain.memory import ConversationBufferMemory
 from langchain.memory.chat_message_histories import StreamlitChatMessageHistory
-from langchain.chains import LLMChain
-from langchain.prompts import PromptTemplate
+from langchain.chains import ConversationalRetrievalChain
 from langchain.embeddings import VoyageEmbeddings
 from langchain.vectorstores import SupabaseVectorStore
 from langchain.llms.huggingface_pipeline import HuggingFacePipeline
 from st_supabase_connection import SupabaseConnection
 
 msgs = StreamlitChatMessageHistory()
-memory = ConversationBufferMemory(memory_key="history", chat_memory=msgs)
+memory = ConversationBufferMemory(
+    memory_key="history", chat_memory=msgs, return_messages=True
+)
 
 supabase_client = st.connection(
     name="orbgpt",
@@ -18,45 +21,51 @@ supabase_client = st.connection(
     ttl=None,
 )
 
-embeddings = VoyageEmbeddings(model="voyage-01")
-vector_store = SupabaseVectorStore(
-    embedding=embeddings,
-    client=supabase_client,
-    table_name="documents",
-    query_name="match_documents",
-)
 
+@st.cache_resource
+def load_retriever():
+    # load embeddings using VoyageAI and Supabase
+    embeddings = VoyageEmbeddings(model="voyage-01")
+    vector_store = SupabaseVectorStore(
+        embedding=embeddings,
+        client=supabase_client.client,
+        table_name="documents",
+        query_name="match_documents",
+    )
+    return vector_store.as_retriever()
 
-model_path = "01-ai/Yi-6B-Chat"
-tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
-model = AutoModelForCausalLM.from_pretrained(
-    model_path,
-    device_map="auto",
-    offload_folder="offload",
-    offload_state_dict=True,
-    torch_dtype="auto",
-).eval()
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=10,
-    use_fast=False,
-)
-hf = HuggingFacePipeline(pipeline=pipe)
 
-template = """Question: {question}
+@st.cache_resource
+def load_model():
+    model_path = "llmware/bling-falcon-1b-0.1"
+    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_path,
+        device_map="auto",
+        offload_folder="offload",
+        offload_state_dict=True,
+        torch_dtype="auto",
+    ).eval()
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        use_fast=False,
+    )
+    return HuggingFacePipeline(pipeline=pipe)
 
-Answer: Let's think step by step."""
-prompt = PromptTemplate.from_template(template)
 
-chain = prompt | hf
+hf = load_model()
+retriever = load_retriever()
+chat = ConversationalRetrievalChain.from_llm(hf, retriever)
 
-question = "What is electroencephalography?"
-
-st.text(chain.invoke({"question": question}))
+st.markdown(
+    "<div style='display: flex;justify-content: center;'><img width='150' src='data:image/png;base64,{}' class='img-fluid'></div>".format(
+        base64.b64encode(Path("orbgptlogo.png").read_bytes()).decode()
+    ),
+    unsafe_allow_html=True,
+)
 
-st.title("🪩🤖")
 
 if len(msgs.messages) == 0:
     msgs.add_ai_message("Ask me anything about orb community projects!")
@@ -66,5 +75,9 @@ for msg in msgs.messages:
 
 if prompt := st.chat_input("Ask something"):
     st.chat_message("human").write(prompt)
-    # Run
-    st.chat_message("ai").write("hehe")
+    msgs.add_user_message(prompt)
+    with st.chat_message("ai"):
+        with st.spinner("Processing your question..."):
+            response = chat({"question": prompt, "chat_history": memory.buffer})
+            msgs.add_ai_message(response["answer"])
+            st.write(response["answer"])