CosmoAI commited on
Commit
7267935
·
1 Parent(s): 19e01ac

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -2
app.py CHANGED
@@ -9,6 +9,10 @@ from langchain.vectorstores import Chroma
9
  from langchain.llms import HuggingFacePipeline
10
  from langchain.chains import RetrievalQA
11
  from constants import CHROMA_SETTINGS
 
 
 
 
12
 
13
  #model and tokenizer loading
14
  checkpoint = "LaMini-T5-738M"
@@ -32,8 +36,10 @@ def llm_pipeline():
32
  @st.cache_resource
33
  def qa_llm():
34
  llm = llm_pipeline()
35
- embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
36
- db = Chroma(persist_directory="db", embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
 
 
37
  retriever = db.as_retriever()
38
  qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents=True)
39
  return qa
 
9
  from langchain.llms import HuggingFacePipeline
10
  from langchain.chains import RetrievalQA
11
  from constants import CHROMA_SETTINGS
12
+ import chromadb
13
+ from chromadb.utils import embedding_functions
14
+
15
+
16
 
17
  #model and tokenizer loading
18
  checkpoint = "LaMini-T5-738M"
 
36
  @st.cache_resource
37
  def qa_llm():
38
  llm = llm_pipeline()
39
+ embeddings = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-MiniLM-L6-v2")
40
+ client = chromadb.PersistentClient(path="/db")
41
+ clientsetting = chromadb.Client(CHROMA_SETTINGS)
42
+ db = Chroma(persist_directory=client, embedding_function=embeddings, client_settings=clientsetting)
43
  retriever = db.as_retriever()
44
  qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=retriever, return_source_documents=True)
45
  return qa