Merge branch 'feat/extract_keyword' into main

sorbobotapp/app.py

@@ -3,12 +3,11 @@ import os
 
 import streamlit as st
 import streamlit.components.v1 as components
-from langchain.callbacks import get_openai_callback
-
 from chain import get_chain
 from chat_history import insert_chat_history, insert_chat_history_articles
 from connection import connect
 from css import load_css
+from langchain.callbacks import get_openai_callback
 from message import Message
 
 st.set_page_config(layout="wide")

sorbobotapp/chain.py

@@ -1,11 +1,10 @@
 import os
 
 import sqlalchemy
+from conversation_retrieval_chain import CustomConversationalRetrievalChain
 from langchain.chains.conversation.memory import ConversationBufferMemory
 from langchain.embeddings import GPT4AllEmbeddings
 from langchain.llms import OpenAI
-
-from conversation_retrieval_chain import CustomConversationalRetrievalChain
 from vector_store import CustomVectorStore
 
 

sorbobotapp/conversation_retrieval_chain.py

@@ -1,12 +1,17 @@
 import inspect
+import json
 from typing import Any, Dict, Optional
 
+from keyword_extraction import KeywordExtractor
 from langchain.callbacks.manager import CallbackManagerForChainRun
 from langchain.chains.conversational_retrieval.base import (
     ConversationalRetrievalChain, _get_chat_history)
+from langchain.schema import Document
 
 
 class CustomConversationalRetrievalChain(ConversationalRetrievalChain):
+    keyword_extractor: KeywordExtractor = KeywordExtractor()
+
     def _handle_docs(self, docs):
         if len(docs) == 0:
             return False, "No documents found. Can you rephrase ?"
@@ -16,6 +21,33 @@ class CustomConversationalRetrievalChain(ConversationalRetrievalChain):
             return False, "Too many documents found. Can you specify your request ?"
         return True, ""
 
+    def rerank_documents(self, question: str, docs: list[Document]) -> list[Document]:
+        """Rerank documents based on the number of similar keywords
+
+        Args:
+            question (str): Orinal question
+            docs (list[Document]): List of documents
+
+        Returns:
+            list[Document]: List of documents sorted by the number of similar keywords
+        """
+        keywords = self.keyword_extractor(question)
+
+        for doc in docs:
+            doc.metadata["similar_keyword"] = 0
+            doc_keywords = json.loads(doc.page_content)["keywords"]
+            if doc_keywords is None:
+                continue
+            doc_keywords = doc_keywords.lower().split(",")
+
+            for kw in keywords:
+                if kw.lower() in doc_keywords:
+                    doc.metadata["similar_keyword"] += 1
+                    print("similar keyword : ", kw)
+
+        docs = sorted(docs, key=lambda x: x.metadata["similar_keyword"])
+        return docs
+
     def _call(
         self,
         inputs: Dict[str, Any],
@@ -40,6 +72,7 @@ class CustomConversationalRetrievalChain(ConversationalRetrievalChain):
             docs = self._get_docs(new_question, inputs, run_manager=_run_manager)
         else:
             docs = self._get_docs(new_question, inputs)  # type: ignore[call-arg]
+
         valid_docs, message = self._handle_docs(docs)
         if not valid_docs:
             return {
@@ -47,6 +80,9 @@ class CustomConversationalRetrievalChain(ConversationalRetrievalChain):
                 "source_documents": docs,
             }
 
+        # Add reranking
+        docs = self.rerank_documents(new_question, docs)
+
         new_inputs = inputs.copy()
         if self.rephrase_question:
             new_inputs["question"] = new_question

sorbobotapp/keyword_extraction.py

@@ -0,0 +1,58 @@
+from typing import Any
+
+from langchain.chat_models import ChatOpenAI
+from langchain.output_parsers import NumberedListOutputParser
+from langchain.prompts import ChatPromptTemplate
+from utils import str_to_list
+
+query_template = """
+You are a bi-lingual (french and english) linguistic teacher working at a top-tier university.
+We are conducting a research project that requires the extraction of keywords from chatbot queries.
+Below, you will find a query. Please identify and rank the three most important keywords or phrases (n-grams) based on their relevance to the main topic of the query.
+For each keyword or phrase, assign it to one of the following categories: ["University / Company", "Research domain", "Country", "Name", "Other"].
+An 'n-gram' refers to a contiguous sequence of words, where 'n' can be 1 for a single word, 2 for a pair of words, and so on, up to two words in length.
+Please ensure not to list more than three n-grams in total.
+Your expertise in linguistic analysis is crucial for the success of this project. Thank you for your contribution.
+
+Please attach your ranked list in the following format:
+1. Keyword/Phrase - Category
+2. Keyword/Phrase - Category
+3. Keyword/Phrase - Category
+
+You must be concise and don't need to justify your choices.
+```
+{query}
+```
+"""
+
+output_parser = NumberedListOutputParser()
+format_instructions = output_parser.get_format_instructions()
+
+
+class KeywordExtractor:
+    def __init__(self):
+        super().__init__()
+        self.model = ChatOpenAI()
+        self.prompt = ChatPromptTemplate.from_template(
+            template=query_template,
+        )
+
+        self.chain = self.prompt | self.model  # | output_parser
+
+    def __call__(
+        self, inputs: str, filter_categories: list[str] = ["Research domain"]
+    ) -> Any:
+        output = self.chain.invoke({"query": inputs})
+
+        keywords = output_parser.parse(output.content)
+
+        filtered_keywords = []
+        for keyword in keywords:
+            if " - " not in keyword:
+                continue
+
+            keyword, category = keyword.split(" - ", maxsplit=2)
+            if category in filter_categories:
+                filtered_keywords.append(keyword)
+
+        return filtered_keywords

sorbobotapp/vector_store.py

@@ -10,11 +10,10 @@ import sqlalchemy
 from langchain.docstore.document import Document
 from langchain.schema.embeddings import Embeddings
 from langchain.vectorstores.base import VectorStore
-from sqlalchemy import delete, text
-from sqlalchemy.orm import Session
-
 from models.article import Article
 from models.distance import DistanceStrategy, distance_strategy_limit
+from sqlalchemy import delete, text
+from sqlalchemy.orm import Session
 from utils import str_to_list
 
 DEFAULT_DISTANCE_STRATEGY = DistanceStrategy.COSINE
@@ -245,7 +244,7 @@ class CustomVectorStore(VectorStore):
                         "doi": result["doi"],
                         "hal_id": result["hal_id"],
                         "distance": result["distance"],
-                        "abstract": result["abstract"],
+                        "abstract": result["abstract"][0],
                     },
                 ),
                 result["distance"] if self.embedding_function is not None else None,