miniondenis/Doc_eater

1from typing import Dict, List2from typing_extensions import TypedDict3 4from langchain_core.documents import Document5 6from lib.model_builder import ModelBuilderV27from lib.vectorestores import FAISSBuilder8from lib.model_builder import ModelBuilderV29from lib.vectorestores import FAISSBuilder10from langchain.retrievers.multi_query import MultiQueryRetriever11from lib.runnables import (12    casual_llm,13    retrieval_grader_3,14    rag_chain,15    message_classificator,16)17from lib.config import Config18from langgraph.graph import END, StateGraph19from transformers import AutoModel, AutoTokenizer20import torch21import torch.nn.functional as F22 23config = Config("config.yml")24model_name = config.get("embeddings", "intfloat", "model")25device = config.get("embeddings", "intfloat", "device")26 27tokenizer = AutoTokenizer.from_pretrained(model_name)28model = AutoModel.from_pretrained(model_name)29 30device = torch.device(device)31model.to(device)32SIMILARITY_TRESHHOLD = 0.833 34 35def get_embeddings(texts):36    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True)37    inputs.to(device)38    with torch.no_grad():39        outputs = model(**inputs)40        embeddings = torch.mean(outputs.last_hidden_state, dim=1)41    return embeddings42 43 44class GraphState(TypedDict):45    """46    Represents the state of our graph.47 48    Attributes:49        question: question50        generation: LLM generation51        web_search: whether to add search52        documents: list of documents53    """54 55    question: str56    generation: str57    documents: List[Document]58    filtered_documets: List[Document]59    count_regenerations: int60 61 62def combine_vectors(vectors):63    result = []64    vec1_count = len(vectors["vector1"])65    # vec2_count = len(vectors["vector2"])66    for i in range(vec1_count):67        if i < vec1_count:68            result.append(vectors["vector1"][i])69        # if i < vec2_count:70        #     result.append(vectors['vector2'][i])71    return result72 73 74def retrieve(state):75    """76    Retrieve documents77 78    Args:79        state (dict): The current graph state80 81    Returns:82        state (dict): New key added to state, documents, that contains retrieved documents83    """84    print("---RETRIEVE---")85    question = state["question"]86 87    # Retrieval88    with FAISSBuilder() as faiss_retriever:89        with ModelBuilderV2("openchat/openchat-7b") as mq_llm:90            retriever = MultiQueryRetriever.from_llm(91                retriever=faiss_retriever, llm=mq_llm92            )93            documents = retriever.get_relevant_documents(question)94            return {"documents": documents, "question": question}95 96 97def start_point(state):98    """99    Start point, just return state100 101    Args:102        state (dict): The current graph state103 104    Returns:105        state (dict): The current graph state106    """107    return state108 109 110def casual_chat(state):111    """112    Define type of message113 114    Args:115        state (dict): The current graph state116 117    Returns:118        state (dict): New key added to state, generation, that contains message with casual answer119    """120    question = state["question"]121    print("---CASUAL CHAT---")122    generation = casual_llm.invoke(123        {"question": question},124        config={125            "configurable": {126                "conversation_id": "default_session",127                "user_id": "deafault_user",128            }129        },130    )131    state["generation"] = generation132 133    return state134 135 136def define_message_type(state):137    """138    Define type of message139 140    Args:141        state (dict): The current graph state142 143    Returns:144 145    """146    print("---MESSAGE CLASSIFICATION---")147    question = state["question"]148    msg_type_obj = message_classificator.invoke({"question": question})149    print(150        f"---MESSAGE TYPE: {msg_type_obj['message_type']} SYSTEM MESSAGE---\n {msg_type_obj['system_message']}"151    )152    msg_type = msg_type_obj["message_type"]153 154    if msg_type == "TAX":155        return "retrieve"156    # if msg_type == "":157    return "casual_chat"158    return "__end__"159 160 161def generate(state):162    """163    Generate answer164 165    Args:166        state (dict): The current graph state167 168    Returns:169        state (dict): New key added to state, generation, that contains LLM generation, based on documents170    """171    print("---GENERATE---")172    question = state["question"]173    documents = state["documents"]174 175    # RAG generation176    generation = rag_chain.invoke(177        {"context": documents, "question": question},178        config={179            "configurable": {180                "conversation_id": "default_session",181                "user_id": "deafault_user",182            }183        },184    )185    return {"documents": documents, "question": question, "generation": generation}186 187 188def grade_documents_by_embed(state):189    """190    Determines whether the retrieved documents are relevant to the question.191 192    Args:193        state (dict): The current graph state194 195    Returns:196        state (dict): Updates documents key with only filtered relevant documents197    """198    question = state["question"]199    documents = state["documents"]200 201    # Score each doc202    filtered_docs = []203 204    query_embedding = get_embeddings([question])205    document_embeddings = get_embeddings([doc.page_content for doc in documents])206 207    # Calculate cosine similarity208    similarity_scores = F.cosine_similarity(query_embedding, document_embeddings)209    for doc, score in zip(documents, similarity_scores):210        if score >= SIMILARITY_TRESHHOLD:211            filtered_docs.append(doc)212    sorted_documents = [213        doc[0]214        for doc in sorted(215            zip(documents, similarity_scores), key=lambda x: x[1], reverse=True216        )217    ]218    cut_off_documents = sorted_documents[:5]219    return {"documents": cut_off_documents, "question": question}220 221 222def grade_documents(state):223    """224    Determines whether the retrieved documents are relevant to the question.225 226    Args:227        state (dict): The current graph state228 229    Returns:230        state (dict): Updates documents key with only filtered relevant documents231    """232 233    print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")234    question = state["question"]235    documents = state["documents"]236 237    # Score each doc238    filtered_docs = []239    count_docs = len(documents)240    for ind_d in range(0, count_docs, 3):241        d_1 = documents[ind_d] if ind_d < count_docs else None242        d_2 = documents[ind_d + 1] if ind_d + 1 < count_docs else None243        d_3 = documents[ind_d + 2] if ind_d + 2 < count_docs else None244        scores = retrieval_grader_3.invoke(245            {246                "question": question,247                "document_1": d_1,248                "document_2": d_2,249                "document_3": d_3,250            }251        )252        for j in range(len(scores)):253            grade = scores[j]["score"]254            if grade > 0.7:255                print(f"---GRADE: DOCUMENT RELEVANT--- GRADE: {grade}")256                filtered_docs.append(documents[ind_d + j])257            else:258                print("---GRADE: DOCUMENT NOT RELEVANT---")259 260    return {"documents": filtered_docs, "question": question}261 262 263def make_collapsable_source_message(doc: Dict):264    file_path = doc.metadata.get("file_name", "")265    file_name = file_path.replace(".pdf", "")266    chapter_title = doc.metadata.get("chapter_title", None)267    page_num = doc.metadata.get("first_page_num", None)268    title = f"""269        {file_name}270        {f": {chapter_title} " if chapter_title is not None else ""}271        {f"Стр. {page_num} " if page_num is not None else ""}272    """.replace(273        "\n", " "274    )275    content = doc.page_content.replace("\n\n", "\n")276 277    if page_num is None:278        message = rf"""279            <details>280            <summary>{title}</summary>281            {str(content)}282            </details>283        """284    else:285        base_url = "http://localhost:5000/sta"286        url = f"{base_url}?file={file_path}&#page={page_num}&zoom=90&toolbar=0"287        message = f"""288            <a class="open_pdf" href='{url}' onclick="return openPdf('{url}')">{title}</a>289        """290 291    return message292 293 294def add_sources(state):295    """296    Determines whether the retrieved documents are relevant to the question.297 298    Args:299        state (dict): The current graph state300 301    Returns:302        state (dict): Add collapsable sources303    """304    question = state["question"]305    documents = state["documents"]306    generation = state["generation"]307 308    sources_message = "<i></i>".join(map(make_collapsable_source_message, documents))309    extended_generation_message = f"{generation} {sources_message}"310    return {311        "documents": documents,312        "question": question,313        "generation": extended_generation_message,314    }315 316 317def build_workflow():318    workflow = StateGraph(GraphState)319 320    # Define the nodes321    workflow.add_node("start_point", start_point)322    workflow.add_node("retrieve", retrieve)  # retrieve323    workflow.add_node("grade_documents", grade_documents_by_embed)  # grade documents324    workflow.add_node("generate", generate)  # generate325    workflow.add_node("casual_chat", casual_chat)  # simple chat326    workflow.add_node("add_sources", add_sources)327    # Build graph328    workflow.set_entry_point("start_point")329    workflow.add_conditional_edges("start_point", define_message_type)330    workflow.add_edge("retrieve", "grade_documents")331    workflow.add_edge("grade_documents", "generate")332    workflow.add_edge("generate", "add_sources")333    workflow.add_edge("add_sources", END)334    workflow.add_edge("casual_chat", END)335    return workflow.compile()336