Spaces:
Sleeping
Sleeping
File size: 3,314 Bytes
278ff72 |
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 |
import os
import tempfile
from chainlit.types import AskFileResponse
from langchain_community.document_loaders import PyMuPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams
from langchain_openai.embeddings import OpenAIEmbeddings
from langchain.storage import LocalFileStore
from langchain_qdrant import QdrantVectorStore
from langchain.embeddings import CacheBackedEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.globals import set_llm_cache
from langchain_openai import ChatOpenAI
from langchain_core.caches import InMemoryCache
from langchain_core.runnables.passthrough import RunnablePassthrough
from uuid import uuid4
from utilities.prompts import get_system_template, get_user_template
def load_file(file: AskFileResponse, chunk_size=1000, chunk_overlap=100):
import tempfile
with tempfile.NamedTemporaryFile(mode="wb", delete=False) as tempfile:
with open(tempfile.name, "wb") as f:
f.write(file.content)
Loader = PyMuPDFLoader
loader = Loader(tempfile.name)
documents = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
docs = text_splitter.split_documents(documents)
for i, doc in enumerate(docs):
doc.metadata["source"] = f"source_{i}"
return docs
def process_embeddings(docs):
core_embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
collection_name = f"pdf_to_parse_{uuid4()}"
client = QdrantClient(":memory:")
client.create_collection(
collection_name=collection_name,
vectors_config=VectorParams(size=1536, distance=Distance.COSINE),
)
# Adding cache!
store = LocalFileStore("./cache/")
cached_embedder = CacheBackedEmbeddings.from_bytes_store(
core_embeddings, store, namespace=core_embeddings.model
)
# Typical QDrant Vector Store Set-up
vectorstore = QdrantVectorStore(
client=client,
collection_name=collection_name,
embedding=cached_embedder)
vectorstore.add_documents(docs)
retriever = vectorstore.as_retriever(search_type="mmr", search_kwargs={"k": 3})
return retriever
def prepare_rag_chain(retriever, prompt_cache="yes"):
print(prompt_cache)
system_template = get_system_template()
user_template = get_user_template()
chat_prompt = ChatPromptTemplate.from_messages([
("system", system_template),
("human", user_template)
])
chat_model = ChatOpenAI(model="gpt-4o-mini")
if prompt_cache == "yes":
set_llm_cache(InMemoryCache())
from operator import itemgetter
rag_qa_chain = (
{"context": itemgetter("question") | retriever, "question": itemgetter("question"), "language": itemgetter("language")}
| RunnablePassthrough.assign(context=itemgetter("context"), language=itemgetter("language"))
| chat_prompt | chat_model
)
return rag_qa_chain
def process_file(file, prompt_cache):
docs = load_file(file)
retriever = process_embeddings(docs)
rag_chain = prepare_rag_chain(retriever, prompt_cache)
return {"chain": rag_chain, "retriever": retriever} |