Spaces:
Sleeping
Sleeping
File size: 2,221 Bytes
47ad957 |
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 |
from langchain.callbacks import FileCallbackHandler
from langchain_community.chat_models import ChatOllama
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from loguru import logger
from rag_101.retriever import (
RAGException,
create_parent_retriever,
load_embedding_model,
load_pdf,
load_reranker_model,
retrieve_context,
)
class RAGClient:
embedding_model = load_embedding_model()
reranker_model = load_reranker_model()
def __init__(self, files, model="mistral"):
docs = load_pdf(files=files)
self.retriever = create_parent_retriever(docs, self.embedding_model)
llm = ChatOllama(model=model)
prompt_template = ChatPromptTemplate.from_template(
(
"Please answer the following question based on the provided `context` that follows the question.\n"
"Think step by step before coming to answer. If you do not know the answer then just say 'I do not know'\n"
"question: {question}\n"
"context: ```{context}```\n"
)
)
self.chain = prompt_template | llm | StrOutputParser()
def stream(self, query: str) -> dict:
try:
context, similarity_score = self.retrieve_context(query)[0]
context = context.page_content
if similarity_score < 0.005:
context = "This context is not confident. " + context
except RAGException as e:
context, similarity_score = e.args[0], 0
logger.info(context)
for r in self.chain.stream({"context": context, "question": query}):
yield r
def retrieve_context(self, query: str):
return retrieve_context(
query, retriever=self.retriever, reranker_model=self.reranker_model
)
def generate(self, query: str) -> dict:
contexts = self.retrieve_context(query)
return {
"contexts": contexts,
"response": self.chain.invoke(
{"context": contexts[0][0].page_content, "question": query}
),
} |