Spaces:
Sleeping
Sleeping
File size: 4,532 Bytes
14c0e58 c4eb0c2 14c0e58 c4eb0c2 14c0e58 c4eb0c2 14c0e58 c4eb0c2 14c0e58 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 |
import os
from operator import itemgetter
import chainlit as cl
import tiktoken
from dotenv import load_dotenv
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.retrievers import MultiQueryRetriever
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_community.document_loaders import PyMuPDFLoader, PythonLoader, NotebookLoader
from langchain_community.vectorstores import Qdrant
from langchain_openai import ChatOpenAI
from langchain_openai.embeddings import OpenAIEmbeddings
# Load environment variables
load_dotenv()
# Configuration for OpenAI
OPENAI_API_KEY = os.environ["OPENAI_API_KEY"]
openai_chat_model = ChatOpenAI(model="gpt-4-turbo", temperature=0.1)
# Define the RAG prompt
RAG_PROMPT = """
CONTEXT:
{context}
QUERY:
{question}
Answer the query in a pretty format if the context is related to it; otherwise, answer: 'Sorry, I can't answer.'
"""
rag_prompt = ChatPromptTemplate.from_template(RAG_PROMPT)
# ChainLit setup for chat interaction
@cl.on_chat_start
async def start_chat():
settings = {
"model": "gpt-3.5-turbo",
"temperature": 0,
"top_p": 1,
"frequency_penalty": 0,
"presence_penalty": 0,
}
cl.user_session.set("settings", settings)
# Display a welcoming message with instructions
welcome_message = "Welcome to the AIMS-Tutor! Please upload a Jupyter notebook (.ipynb and max. 5mb) to start."
await cl.Message(content=welcome_message).send()
# Wait for the user to upload a file
files = None
while files is None:
files = await cl.AskFileMessage(
content="Please upload a Jupyter notebook (.ipynb, max. 5mb):",
accept={"application/x-ipynb+json": [".ipynb"]},
max_size_mb=5
).send()
file = files[0] # Get the first file
if file:
# Load the Jupyter notebook
notebook_path = file.path # Extract the path from the AskFileResponse object
loader = NotebookLoader(
notebook_path,
include_outputs=False,
max_output_length=20,
remove_newline=True,
traceback=False
)
docs = loader.load()
cl.user_session.set("docs", docs) # Store the docs in the user session
# Initialize the retriever components after loading document
text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=50, length_function=tiktoken_len) # Initialize the text splitter
split_chunks = text_splitter.split_documents(docs) # Split the documents into chunks
embedding_model = OpenAIEmbeddings(model="text-embedding-3-small") # Initialize the embedding model
qdrant_vectorstore = Qdrant.from_documents(split_chunks, embedding_model, location=":memory:", collection_name="Notebook") # Create a Qdrant vector store
qdrant_retriever = qdrant_vectorstore.as_retriever() # Set the Qdrant vector store as a retriever
multiquery_retriever = MultiQueryRetriever.from_llm(retriever=qdrant_retriever, llm=openai_chat_model, include_original=True) # Create a multi-query retriever on top of the Qdrant retriever
# Store the multiquery_retriever in the user session
cl.user_session.set("multiquery_retriever", multiquery_retriever)
@cl.on_message
async def main(message: cl.Message):
# Retrieve the multi-query retriever from session
multiquery_retriever = cl.user_session.get("multiquery_retriever")
if not multiquery_retriever:
await cl.Message(content="No document processing setup found. Please upload a Jupyter notebook first.").send()
return
question = message.content
response = handle_query(question, multiquery_retriever) # Process the question
msg = cl.Message(content=response)
await msg.send()
def handle_query(question, retriever):
# Define the retrieval augmented query-answering chain
retrieval_augmented_qa_chain = (
{"context": itemgetter("question") | retriever, "question": itemgetter("question")}
| RunnablePassthrough.assign(context=itemgetter("context"))
| {"response": rag_prompt | openai_chat_model, "context": itemgetter("context")}
)
response = retrieval_augmented_qa_chain.invoke({"question": question})
return response["response"].content
# Tokenization function
def tiktoken_len(text):
tokens = tiktoken.encoding_for_model("gpt-3.5-turbo").encode(text)
return len(tokens)
|