Cheselle commited on
Commit
b83d3fb
·
verified ·
1 Parent(s): 41b1d44

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -14
app.py CHANGED
@@ -3,22 +3,19 @@ from langchain.prompts import ChatPromptTemplate
3
  from langchain.schema import StrOutputParser
4
  from langchain.schema.runnable import Runnable
5
  from langchain.schema.runnable.config import RunnableConfig
6
- from typing import cast
7
  from dotenv import load_dotenv
8
  import os
9
  from langchain_community.document_loaders import PyMuPDFLoader
10
  from langchain.text_splitter import RecursiveCharacterTextSplitter
11
- from langchain_openai.embeddings import OpenAIEmbeddings
12
  from langchain_community.vectorstores import Qdrant
13
- from langchain_core.runnables import RunnablePassthrough, RunnableParallel
14
  import chainlit as cl
15
- from pathlib import Path
16
  from sentence_transformers import SentenceTransformer
17
 
 
18
  load_dotenv()
19
-
20
  os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
21
 
 
22
  class SentenceTransformerEmbedding:
23
  def __init__(self, model_name):
24
  self.model = SentenceTransformer(model_name)
@@ -29,7 +26,7 @@ class SentenceTransformerEmbedding:
29
  def __call__(self, texts):
30
  return self.embed_documents(texts) # Make it callable
31
 
32
- @cl.on_chat_start
33
  async def on_chat_start():
34
  model = ChatOpenAI(streaming=True)
35
 
@@ -62,20 +59,19 @@ async def on_chat_start():
62
  sentence_blueprint = metadata_generator(ai_blueprint_document, "AI Blueprint", sentence_text_splitter)
63
 
64
  sentence_combined_documents = sentence_framework + sentence_blueprint
65
- from langchain_community.embeddings import HuggingFaceEmbeddings # Use HuggingFaceEmbeddings if needed
66
-
67
 
68
- # Initialize the custom embedding class
69
  embedding_model = SentenceTransformerEmbedding('Cheselle/finetuned-arctic-sentence')
70
 
71
- # Create the Qdrant vector store using the custom embedding model
72
  sentence_vectorstore = Qdrant.from_documents(
73
  documents=sentence_combined_documents,
74
- embedding_function=embedding_model, # Ensure this is an instance
75
  location=":memory:",
76
  collection_name="AI Policy"
77
  )
78
 
 
79
  sentence_retriever = sentence_vectorstore.as_retriever()
80
 
81
  # Check if retriever is initialized correctly
@@ -87,11 +83,10 @@ async def on_chat_start():
87
  cl.user_session.set("retriever", sentence_retriever)
88
  cl.user_session.set("prompt_template", rag_prompt)
89
 
90
-
91
- @cl.on_message
92
  async def on_message(message: cl.Message):
93
  # Get the stored model, retriever, and prompt
94
- model = cast(ChatOpenAI, cl.user_session.get("runnable"))
95
  retriever = cl.user_session.get("retriever")
96
  prompt_template = cl.user_session.get("prompt_template")
97
 
 
3
  from langchain.schema import StrOutputParser
4
  from langchain.schema.runnable import Runnable
5
  from langchain.schema.runnable.config import RunnableConfig
 
6
  from dotenv import load_dotenv
7
  import os
8
  from langchain_community.document_loaders import PyMuPDFLoader
9
  from langchain.text_splitter import RecursiveCharacterTextSplitter
 
10
  from langchain_community.vectorstores import Qdrant
 
11
  import chainlit as cl
 
12
  from sentence_transformers import SentenceTransformer
13
 
14
+ # Load environment variables
15
  load_dotenv()
 
16
  os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
17
 
18
+ # Custom embedding class for SentenceTransformer
19
  class SentenceTransformerEmbedding:
20
  def __init__(self, model_name):
21
  self.model = SentenceTransformer(model_name)
 
26
  def __call__(self, texts):
27
  return self.embed_documents(texts) # Make it callable
28
 
29
+ @cl.on_chat_start # Marks the function to be executed at the start of a user session
30
  async def on_chat_start():
31
  model = ChatOpenAI(streaming=True)
32
 
 
59
  sentence_blueprint = metadata_generator(ai_blueprint_document, "AI Blueprint", sentence_text_splitter)
60
 
61
  sentence_combined_documents = sentence_framework + sentence_blueprint
 
 
62
 
63
+ # Initialize the embedding model instance
64
  embedding_model = SentenceTransformerEmbedding('Cheselle/finetuned-arctic-sentence')
65
 
66
+ # Create the Qdrant vector store using the embedding instance
67
  sentence_vectorstore = Qdrant.from_documents(
68
  documents=sentence_combined_documents,
69
+ embedding=embedding_model, # Pass the embedding instance correctly
70
  location=":memory:",
71
  collection_name="AI Policy"
72
  )
73
 
74
+ # Create retriever from the vector store
75
  sentence_retriever = sentence_vectorstore.as_retriever()
76
 
77
  # Check if retriever is initialized correctly
 
83
  cl.user_session.set("retriever", sentence_retriever)
84
  cl.user_session.set("prompt_template", rag_prompt)
85
 
86
+ @cl.on_message # Marks a function to run each time a message is received
 
87
  async def on_message(message: cl.Message):
88
  # Get the stored model, retriever, and prompt
89
+ model = cl.user_session.get("runnable")
90
  retriever = cl.user_session.get("retriever")
91
  prompt_template = cl.user_session.get("prompt_template")
92