leo-bourrel commited on
Commit
4bcb630
·
1 Parent(s): 503aad2

feat : use custom pg vector and gpt4all embeddings

Browse files
Files changed (1) hide show
  1. app.py +9 -11
app.py CHANGED
@@ -1,18 +1,15 @@
1
- import os
2
-
3
  import streamlit as st
4
  import streamlit.components.v1 as components
5
  from css import load_css
 
6
  from langchain import OpenAI
7
  from langchain.callbacks import get_openai_callback
8
  from langchain.chains import ConversationalRetrievalChain
9
  from langchain.chains.conversation.memory import ConversationBufferMemory
10
- from langchain.embeddings.openai import OpenAIEmbeddings
11
- from langchain.vectorstores.pgvector import PGVector
12
  from message import Message
13
 
14
  CONNECTION_STRING = "postgresql+psycopg2://localhost/sorbobot"
15
- COLLECTION_NAME = ""
16
 
17
 
18
  def initialize_session_state():
@@ -21,15 +18,16 @@ def initialize_session_state():
21
  if "token_count" not in st.session_state:
22
  st.session_state.token_count = 0
23
  if "conversation" not in st.session_state:
24
- embeddings = OpenAIEmbeddings()
25
 
26
- store = PGVector(
27
- collection_name=COLLECTION_NAME,
28
- connection_string=CONNECTION_STRING,
29
  embedding_function=embeddings,
 
 
 
30
  )
31
 
32
- retriever = store.as_retriever()
33
 
34
  llm = OpenAI(
35
  temperature=0,
@@ -39,7 +37,7 @@ def initialize_session_state():
39
 
40
  st.session_state.memory = ConversationBufferMemory()
41
  st.session_state.conversation = ConversationalRetrievalChain.from_llm(
42
- llm=llm, retriever=retriever
43
  )
44
 
45
 
 
 
 
1
  import streamlit as st
2
  import streamlit.components.v1 as components
3
  from css import load_css
4
+ from custom_pgvector import CustomPGVector
5
  from langchain import OpenAI
6
  from langchain.callbacks import get_openai_callback
7
  from langchain.chains import ConversationalRetrievalChain
8
  from langchain.chains.conversation.memory import ConversationBufferMemory
9
+ from langchain.embeddings import GPT4AllEmbeddings
 
10
  from message import Message
11
 
12
  CONNECTION_STRING = "postgresql+psycopg2://localhost/sorbobot"
 
13
 
14
 
15
  def initialize_session_state():
 
18
  if "token_count" not in st.session_state:
19
  st.session_state.token_count = 0
20
  if "conversation" not in st.session_state:
21
+ embeddings = GPT4AllEmbeddings()
22
 
23
+ db = CustomPGVector(
 
 
24
  embedding_function=embeddings,
25
+ table_name="article",
26
+ column_name="abstract_embedding",
27
+ connection_string=CONNECTION_STRING,
28
  )
29
 
30
+ retriever = db.as_retriever()
31
 
32
  llm = OpenAI(
33
  temperature=0,
 
37
 
38
  st.session_state.memory = ConversationBufferMemory()
39
  st.session_state.conversation = ConversationalRetrievalChain.from_llm(
40
+ llm=llm, retriever=retriever, verbose=True
41
  )
42
 
43