upload

Dockerfile

@@ -14,13 +14,18 @@ RUN pip install --no-cache-dir -r requirements.txt
 # Copy the current directory contents into the container at /app
 COPY . /app
 
-# Expose port 7860 for the Flask app to listen on
+# Expose port 7860 for the streamlit app to listen on
 EXPOSE 7860
 
+# open the grobid server
 RUN apt-get update && \
     apt-get install -y wget unzip
 RUN chmod +x serve_grobid.sh
 RUN ./serve_grobid.sh 
 
-# Run the command to start the Flask app
-CMD ["python", "app.py"]
+# open the backend server
+RUN python backend.py &
+
+# Run the command to start the streamlit app
+CMD ["streamlit", "run", "frontend.py", "--server.address", "0.0.0.0", "--server.port", "7860"]
+

LICENSE

@@ -0,0 +1,25 @@
+BSD 2-Clause License
+
+Copyright (c) 2023, Yixin Liu
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+* Redistributions of source code must retain the above copyright notice, this
+  list of conditions and the following disclaimer.
+
+* Redistributions in binary form must reproduce the above copyright notice,
+  this list of conditions and the following disclaimer in the documentation
+  and/or other materials provided with the distribution.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

README.md

@@ -1,11 +1,52 @@
----
-title: ChatPaper
-emoji: 📊
-colorFrom: purple
-colorTo: blue
-sdk: docker
-pinned: false
-license: bsd-2-clause
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# ChatPaper
+Yet another paper reading assistant, similar as [ChatPDF](https://www.chatpdf.com/). 
+
+## Setup
+
+1. Install dependencies (tested on Python 3.9)
+
+```bash
+ pip install -r requirements.txt
+```
+
+2. Setup GROBID local server
+
+```bash
+bash serve_grobid.sh
+```
+
+3. Setup backend
+
+```bash
+python backend.py --port 5000 --host localhost
+```
+
+4. Frontend 
+
+```bash
+streamlit run frontend.py --server.port 8502 --server.host localhost
+```
+
+## Demo Example
+
+- Prepare an [OpenAI API key](https://platform.openai.com/account/api-keys) and then upload a PDF to start chatting with the paper. 
+
+![image-20230318232056584](https://s2.loli.net/2023/03/19/SbsuLQJpdqePoZV.png)
+
+## Implementation Details
+
+- Greedy Dynamic Context: Since the max token limit, we select the most relevant paragraphs in the pdf for each user query. Our model split the text input and output by the chatbot into four part: system_prompt (S), dynamic_source (D), user_query (Q), and model_answer(A). So upon each query, we first rank all the paragraphs by using a sentence_embedding model to calculate the similarity distance between the query embedding and all source embeddings. Then we compose the dynamic_source using a greedy method by to gradually push all relevant paragraphs (maintaing D <= MAX_TOKEN_LIMIT - Q - S - A - SOME_OVERHEAD). 
+
+- Context Truncating: When context is too long, we now we simply pop out the first QA-pair. 
+
+## TODO
+
+- [ ] **Context Condense**: how to deal with long context? maybe we can tune a soft prompt to condense the context
+- [ ] **Poping context out based on similarity**
+
+## References
+
+1. SciPDF Parser: https://github.com/titipata/scipdf_parser 
+2. St-chat: https://github.com/AI-Yash/st-chat
+3. Sentence-transformers: https://github.com/UKPLab/sentence-transformers
+4. ChatGPT Chatbot Wrapper: https://github.com/acheong08/ChatGPT

app.py

@@ -1,10 +0,0 @@
-from flask import Flask
-
-app = Flask(__name__)
-
-@app.route('/')
-def hello_world():
-    return 'Hello, World!'
-
-if __name__ == '__main__':
-    app.run(debug=True, host='0.0.0.0', port=7860)

backend.py

@@ -0,0 +1,81 @@
+from flask import jsonify, Flask, request
+from embedding_model import HuggingfaceSentenceTransformerModel
+from similarity_metric import CosineSimilarity
+from pdf_parser import GrobidSciPDFPaser
+from chatbot import OpenAIChatbot
+from chat_pdf import ChatPDF
+from config import DEFAULT_ENGINE, MAX_TOKEN_MODEL_MAP, DEFAULT_TEMPERATURE, DEFAULT_TOP_P, DEFAULT_PRESENCE_PENALTY, DEFAULT_FREQUENCY_PENALTY, DEFAULT_REPLY_COUNT
+app = Flask(__name__)
+chatpdf_pool = {}
+
+embedding_model = HuggingfaceSentenceTransformerModel()
+simi_metric = CosineSimilarity()
+
+
+@app.route("/query/", methods=['POST', 'GET'])
+def query():
+    api_key = request.headers.get('Api-Key')
+    pdf_link = request.json['pdf_link']
+    user_stamp = request.json['user_stamp']
+    user_query = request.json['user_query']
+    print(
+        "api_key", api_key,
+        "pdf_link", pdf_link,
+        "user_stamp", user_stamp,
+        "user_query", user_query
+    )
+
+    chat_pdf = None
+    if user_stamp not in chatpdf_pool:
+        print(f"User {user_stamp} not in pool, creating new chatpdf")
+        # Initialize the ChatPDF
+        bot = OpenAIChatbot(
+            api_key=api_key,
+            engine=DEFAULT_ENGINE,
+            proxy=None,
+            max_tokens=4000,
+            temperature=DEFAULT_TEMPERATURE,
+            top_p=DEFAULT_TOP_P,
+            presence_penalty=DEFAULT_PRESENCE_PENALTY,
+            frequency_penalty=DEFAULT_FREQUENCY_PENALTY,
+            reply_count=DEFAULT_REPLY_COUNT
+        )
+
+        pdf = GrobidSciPDFPaser(
+            pdf_link=pdf_link
+        )
+        chat_pdf = ChatPDF(
+            pdf=pdf,
+            bot=bot,
+            embedding_model=embedding_model,
+            similarity_metric=simi_metric,
+            user_stamp=user_stamp
+        )
+        chatpdf_pool[user_stamp] = chat_pdf
+    else:
+        print("user_stamp", user_stamp, "already exists")
+        chat_pdf = chatpdf_pool[user_stamp]
+
+    try:
+        response = chat_pdf.chat(user_query)
+        code = 200
+        json_dict = {
+            "code": code,
+            "response": response
+        }
+    except Exception as e:
+        code = 500
+        json_dict = {
+            "code": code,
+            "response": str(e)
+        }
+    return jsonify(json_dict)
+
+
+# @app.route("/", methods=['GET'])
+# def index():
+#     return "Hello World!"
+
+
+if __name__ == '__main__':
+    app.run(host='localhost', port=5000, debug=False)

base_class.py

@@ -0,0 +1,106 @@
+import abc
+import pandas as pd
+import pickle
+
+
+class SimilarityAlg(metaclass=abc.ABCMeta):
+    """Similarity Algorithm to compute similarity between query_embedding and embeddings"""
+
+    def __init__(self) -> None:
+        pass
+
+    @abc.abstractmethod
+    def __call__(self, query_embedding, embeddings) -> None:
+        pass
+
+
+class Embedding_Model(metaclass=abc.ABCMeta):
+    """Embedding Model to compute embedding of a text"""
+
+    def __init__(self, model_name) -> None:
+        """Initialize the embedding model"""
+        embedding_cache_path = f"/data/ckpt/embedding_cache_{model_name}.pkl"
+        self.embedding_cache_path = embedding_cache_path
+
+        # load the cache if it exists, and save a copy to disk
+        try:
+            embedding_cache = pd.read_pickle(embedding_cache_path)
+        except FileNotFoundError:
+            embedding_cache = {}
+        with open(embedding_cache_path, "wb") as embedding_cache_file:
+            pickle.dump(embedding_cache, embedding_cache_file)
+        self.embedding_cache = embedding_cache
+        self.model_name = model_name
+
+    @abc.abstractmethod
+    def __call__(self, text) -> None:
+        """Compute the embedding of the text"""
+        pass
+
+
+class AbstractPDFParser(metaclass=abc.ABCMeta):
+    """ PDF parser to parse a PDF file"""
+
+    def __init__(self, db_name) -> None:
+        """Initialize the pdf database"""
+        db_cache_path = f"/data/ckpt/pdf_parser_{db_name}.pkl"
+        self.db_cache_path = db_cache_path
+
+        # load the cache if it exists, and save a copy to disk
+        try:
+            db_cache = pd.read_pickle(db_cache_path)
+        except FileNotFoundError:
+            db_cache = {}
+        with open(db_cache_path, "wb") as cache_file:
+            pickle.dump(db_cache, cache_file)
+        self.db_cache = db_cache
+        self.db_name = db_name
+
+    @abc.abstractmethod
+    def parse_pdf(self,) -> None:
+        """Parse the PDF file"""
+        pass
+
+    @abc.abstractmethod
+    def _get_metadata(self, ) -> None:
+        """Get the metadata of the PDF file"""
+        pass
+
+    def get_paragraphs(self, ) -> None:
+        """Get the paragraphs of the PDF file"""
+        pass
+
+    @abc.abstractmethod
+    def get_split_paragraphs(self, ) -> None:
+        """
+        Get the split paragraphs of the PDF file
+        Return:
+            split_paragraphs: dict of metadata and corresponding list of split paragraphs
+        """
+        pass
+
+    def _determine_metadata_of_paragraph(self, paragraph) -> None:
+        """
+        Determine the metadata of a paragraph
+        Return:
+            metadata: metadata of the paragraph
+        """
+        pass
+
+    @abc.abstractmethod
+    def _determine_optimal_split_of_pargraphs(self, ) -> None:
+        """
+        Determine the optimal split of paragraphs
+        Return:
+            split_paragraphs: dict of metadata and corresponding list of split paragraphs
+        """
+        pass
+
+
+class ChatbotEngine(metaclass=abc.ABCMeta):
+    def __init__(self,) -> None:
+        pass
+
+    @abc.abstractmethod
+    def query(self, user_query):
+        pass

chat_pdf.py

@@ -0,0 +1,83 @@
+
+class ChatPDF():
+    """ChatPDF enables us to chat with a PDF file
+    """
+
+    def __init__(self, pdf, bot, embedding_model, similarity_metric, expect_answer_token_length_max=100, expect_q_token_length_max=100, user_stamp=None) -> None:
+        self.pdf = pdf
+        self.bot = bot
+        self.embedding_model = embedding_model
+        self.similarity_metric = similarity_metric
+        self.user_stamp = user_stamp
+
+        self.system_task_prompt = f"You are a helpful PDF file. Your task is to provide information and answer any questions related to the topic of {self.pdf.metadata['title']}. You should use the sections of the PDF as your source of information and try to provide concise and accurate answers to any questions asked by the user. If you are unable to find relevant information in the given sections, you will need to let the user know that the source does not contain relevant information but still try to provide an answer based on your general knowledge. You must refer to the corresponding section name and page that you refer to when answering. The following is the related information about the PDF file that will help you answer users' questions:\n\n"
+        self.system_information_prompt = "Title:\n" + self.pdf.metadata['title'] + "\n\nAbstract:\n" + self.pdf.metadata["abstract"] + \
+            "\n\nFiltered paragraphs from each sections (the section titles are enclosed in asterisks**):\n\n"
+
+        self.system_token_length = self.bot.encode_length(
+            self.system_task_prompt) + self.bot.encode_length(self.system_information_prompt)
+        self.expect_answer_token_length_max = expect_answer_token_length_max
+        self.expect_q_token_length_max = expect_q_token_length_max
+        self.context_max_length = self.bot.max_tokens - self.system_token_length - \
+            self.bot.overhead_token - self.expect_answer_token_length_max - \
+            self.expect_q_token_length_max
+
+    def _get_related_context(self, user_query):
+        all_contextes = [user_query]+self.pdf.flattn_paragraphs
+        rank_indices = self.rank_indices(all_contextes, 0,)
+        rank_indices = list(rank_indices)
+        rank_indices.remove(0)
+        inital_context = ":\n\n".join(self.pdf.section_names_with_page_index)
+        context_dict = {section_name: []
+                        for section_name in self.pdf.section_names}
+        inital_context_token_length = self.bot.encode_length(inital_context)
+        running_length = inital_context_token_length
+        for idx in rank_indices:
+            text_to_insert = all_contextes[idx]
+            text_to_insert_token_length = self.bot.encode_length(
+                text_to_insert)
+            if running_length + text_to_insert_token_length < self.context_max_length:
+                running_length += text_to_insert_token_length
+                section = self.pdf.content2section[text_to_insert]
+                context_dict[section].append(text_to_insert)
+            else:
+                break
+        composed_context = ""
+        for i, section_name in enumerate(self.pdf.section_names):
+            if len(context_dict[section_name]) > 0:
+                section_name_with_page_index = self.pdf.section_names_with_page_index[i]
+                composed_context += "**"+section_name_with_page_index + "**" + \
+                    ":\n" + "\n".join(context_dict[section_name]) + "\n\n"
+        return composed_context
+
+    def chat(self, user_query):
+        """Chat with the PDF file
+        """
+        context_data = self._get_related_context(user_query)
+        dynamic_system_context = self.system_task_prompt + \
+            self.system_information_prompt + context_data
+        print(
+            "************ Start of Composed Context ************\n",
+            dynamic_system_context,
+            "\n************ End of Composed Context ************\n"
+        )
+        response = self.bot.query(
+            context=dynamic_system_context, questions=user_query, convo_id=self.user_stamp)
+        return response
+
+    def rank_indices(
+        self,
+        strings: list[str],
+        index_of_source_string: int,
+    ) -> list[int]:
+        """Rank the indices of the strings in the list based on their similarity to the source string."""
+        # get embeddings for all strings
+        embeddings = [self.embedding_model(string) for string in strings]
+        # get the embedding of the source string
+        query_embedding = embeddings[index_of_source_string]
+        # get distances between the source embedding and other embeddings (function from embeddings_utils.py)
+        distances = self.similarity_metric(query_embedding, embeddings)
+        # get rank of indices based on distances
+        import numpy as np
+        indices_of_nearest_neighbors = np.argsort(distances, )
+        return indices_of_nearest_neighbors

chatbot.py

@@ -0,0 +1,320 @@
+from base_class import ChatbotEngine
+import os
+import openai
+import json
+import os
+import requests
+import tiktoken
+from config import MAX_TOKEN_MODEL_MAP
+from utils import get_filtered_keys_from_object
+
+
+class ChatbotWrapper:
+    """
+    Wrapper of Official ChatGPT API,
+    # base on https://github.com/ChatGPT-Hackers/revChatGPT
+    """
+
+    def __init__(
+        self,
+        api_key: str,
+        engine: str = os.environ.get("GPT_ENGINE") or "gpt-3.5-turbo",
+        proxy: str = None,
+        max_tokens: int = 3000,
+        temperature: float = 0.5,
+        top_p: float = 1.0,
+        presence_penalty: float = 0.0,
+        frequency_penalty: float = 0.0,
+        reply_count: int = 1,
+        system_prompt: str = "You are ChatGPT, a large language model trained by OpenAI. Respond conversationally",
+        overhead_token=96,
+    ) -> None:
+        """
+        Initialize Chatbot with API key (from https://platform.openai.com/account/api-keys)
+        """
+        self.engine = engine
+        self.session = requests.Session()
+        self.api_key = api_key
+        self.system_prompt = system_prompt
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+        self.top_p = top_p
+        self.presence_penalty = presence_penalty
+        self.frequency_penalty = frequency_penalty
+        self.reply_count = reply_count
+        self.max_limit = MAX_TOKEN_MODEL_MAP[self.engine]
+        self.overhead_token = overhead_token
+
+        if proxy:
+            self.session.proxies = {
+                "http": proxy,
+                "https": proxy,
+            }
+
+        self.conversation: dict = {
+            "default": [
+                {
+                    "role": "system",
+                    "content": system_prompt,
+                },
+            ],
+        }
+
+        if max_tokens > self.max_limit - self.overhead_token:
+            raise Exception(
+                f"Max tokens cannot be greater than {self.max_limit- self.overhead_token}")
+
+        if self.get_token_count("default") > self.max_tokens:
+            raise Exception("System prompt is too long")
+
+    def add_to_conversation(
+        self,
+        message: str,
+        role: str,
+        convo_id: str = "default",
+    ) -> None:
+        """
+        Add a message to the conversation
+        """
+        self.conversation[convo_id].append({"role": role, "content": message})
+
+    def __truncate_conversation(self, convo_id: str = "default") -> None:
+        """
+        Truncate the conversation
+        """
+        # TODO: context condense with soft prompt tuning
+        while True:
+            if (
+                self.get_token_count(convo_id) > self.max_tokens
+                and len(self.conversation[convo_id]) > 1
+            ):
+                # Don't remove the first message and remove the first QA pair
+                self.conversation[convo_id].pop(1)
+                self.conversation[convo_id].pop(1)
+                # TODO: optimal pop out based on similarity distance
+            else:
+                break
+
+    # https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
+    def get_token_count(self, convo_id: str = "default") -> int:
+        """
+        Get token count
+        """
+        if self.engine not in ["gpt-3.5-turbo", "gpt-3.5-turbo-0301"]:
+            raise NotImplementedError("Unsupported engine {self.engine}")
+
+        encoding = tiktoken.encoding_for_model(self.engine)
+
+        num_tokens = 0
+        for message in self.conversation[convo_id]:
+            # every message follows <im_start>{role/name}\n{content}<im_end>\n
+            num_tokens += 4
+            for key, value in message.items():
+                num_tokens += len(encoding.encode(value))
+                if key == "name":  # if there's a name, the role is omitted
+                    num_tokens += 1  # role is always required and always 1 token
+        num_tokens += 2  # every reply is primed with <im_start>assistant
+        return num_tokens
+
+    def get_max_tokens(self, convo_id: str) -> int:
+        """
+        Get max tokens
+        """
+        return self.max_tokens - self.get_token_count(convo_id)
+
+    def ask_stream(
+        self,
+        prompt: str,
+        role: str = "user",
+        convo_id: str = "default",
+        dynamic_system_prompt=None,
+        **kwargs,
+    ) -> str:
+        """
+        Ask a question
+        """
+        # Make conversation if it doesn't exist
+        if convo_id not in self.conversation:
+            self.reset(convo_id=convo_id, system_prompt=dynamic_system_prompt)
+
+        # adjust system prompt
+        assert dynamic_system_prompt is not None
+        self.conversation[convo_id][0]["content"] = dynamic_system_prompt
+
+        self.add_to_conversation(prompt, "user", convo_id=convo_id)
+        print(" total tokens:")
+        print(self.get_token_count(convo_id))
+        self.__truncate_conversation(convo_id=convo_id)
+        # Get response
+        response = self.session.post(
+            os.environ.get(
+                "API_URL") or "https://api.openai.com/v1/chat/completions",
+            headers={
+                "Authorization": f"Bearer {kwargs.get('api_key', self.api_key)}"},
+            json={
+                "model": self.engine,
+                "messages": self.conversation[convo_id],
+                "stream": True,
+                # kwargs
+                "temperature": kwargs.get("temperature", self.temperature),
+                "top_p": kwargs.get("top_p", self.top_p),
+                "presence_penalty": kwargs.get(
+                    "presence_penalty",
+                    self.presence_penalty,
+                ),
+                "frequency_penalty": kwargs.get(
+                    "frequency_penalty",
+                    self.frequency_penalty,
+                ),
+                "n": kwargs.get("n", self.reply_count),
+                "user": role,
+                "max_tokens": self. get_max_tokens(convo_id=convo_id),
+            },
+            stream=True,
+        )
+        if response.status_code != 200:
+            raise Exception(
+                f"Error: {response.status_code} {response.reason} {response.text}",
+            )
+        response_role: str = None
+        full_response: str = ""
+        for line in response.iter_lines():
+            if not line:
+                continue
+            # Remove "data: "
+            line = line.decode("utf-8")[6:]
+            if line == "[DONE]":
+                break
+            resp: dict = json.loads(line)
+            choices = resp.get("choices")
+            if not choices:
+                continue
+            delta = choices[0].get("delta")
+            if not delta:
+                continue
+            if "role" in delta:
+                response_role = delta["role"]
+            if "content" in delta:
+                content = delta["content"]
+                full_response += content
+                yield content
+        self.add_to_conversation(
+            full_response, response_role, convo_id=convo_id)
+
+    def ask(
+        self,
+        prompt: str,
+        role: str = "user",
+        convo_id: str = "default",
+        dynamic_system_prompt: str = None,
+        **kwargs,
+    ) -> str:
+        """
+        Non-streaming ask
+        """
+        response = self.ask_stream(
+            prompt=prompt,
+            role=role,
+            convo_id=convo_id,
+            dynamic_system_prompt=dynamic_system_prompt,
+            **kwargs,
+        )
+        full_response: str = "".join(response)
+        return full_response
+
+    def rollback(self, n: int = 1, convo_id: str = "default") -> None:
+        """
+        Rollback the conversation
+        """
+        for _ in range(n):
+            self.conversation[convo_id].pop()
+
+    def reset(self, convo_id: str = "default", system_prompt: str = None) -> None:
+        """
+        Reset the conversation
+        """
+        self.conversation[convo_id] = [
+            {"role": "system", "content": system_prompt or self.system_prompt},
+        ]
+
+    def save(self, file: str, *keys: str) -> None:
+        """
+        Save the Chatbot configuration to a JSON file
+        """
+        with open(file, "w", encoding="utf-8") as f:
+            json.dump(
+                {
+                    key: self.__dict__[key]
+                    for key in get_filtered_keys_from_object(self, *keys)
+                },
+                f,
+                indent=2,
+                # saves session.proxies dict as session
+                default=lambda o: o.__dict__["proxies"],
+            )
+
+    def load(self, file: str, *keys: str) -> None:
+        """
+        Load the Chatbot configuration from a JSON file
+        """
+        with open(file, encoding="utf-8") as f:
+            # load json, if session is in keys, load proxies
+            loaded_config = json.load(f)
+            keys = get_filtered_keys_from_object(self, *keys)
+
+            if "session" in keys and loaded_config["session"]:
+                self.session.proxies = loaded_config["session"]
+            keys = keys - {"session"}
+            self.__dict__.update({key: loaded_config[key] for key in keys})
+
+
+class OpenAIChatbot(ChatbotEngine):
+    def __init__(self, api_key: str,
+                 engine: str = os.environ.get("GPT_ENGINE") or "gpt-3.5-turbo",
+                 proxy: str = None,
+                 max_tokens: int = 3000,
+                 temperature: float = 0.5,
+                 top_p: float = 1.0,
+                 presence_penalty: float = 0.0,
+                 frequency_penalty: float = 0.0,
+                 reply_count: int = 1,
+                 system_prompt: str = "You are ChatGPT, a large language model trained by OpenAI. Respond conversationally",
+                 overhead_token=96) -> None:
+        openai.api_key = api_key
+        self.api_key = api_key
+        self.engine = engine
+        self.proxy = proxy
+        self.max_tokens = max_tokens
+        self.temperature = temperature
+        self.top_p = top_p
+        self.presence_penalty = presence_penalty
+        self.frequency_penalty = frequency_penalty
+        self.reply_count = reply_count
+        self.system_prompt = system_prompt
+
+        self.bot = ChatbotWrapper(
+            api_key=self.api_key,
+            engine=self.engine,
+            proxy=self.proxy,
+            max_tokens=self.max_tokens,
+            temperature=self.temperature,
+            top_p=self.top_p,
+            presence_penalty=self.presence_penalty,
+            frequency_penalty=self.frequency_penalty,
+            reply_count=self.reply_count,
+            system_prompt=self.system_prompt,
+            overhead_token=overhead_token
+        )
+        self.overhead_token = overhead_token
+        import tiktoken
+        self.encoding = tiktoken.encoding_for_model(self.engine)
+
+    def encode_length(self, text: str) -> int:
+        return len(self.encoding.encode(text))
+
+    def query(self, questions: str,
+              role: str = "user",
+              convo_id: str = "default",
+              context: str = None,
+              **kwargs,):
+        return self.bot.ask(prompt=questions, role=role, convo_id=convo_id, dynamic_system_prompt=context, **kwargs)

config.py

@@ -0,0 +1,15 @@
+
+
+MAX_TOKEN_MODEL_MAP = {
+    "gpt-3.5-turbo": 4096,
+}
+
+PDF_SAVE_DIR = "./files/"
+
+
+DEFAULT_ENGINE = "gpt-3.5-turbo"
+DEFAULT_TEMPERATURE = 0.5
+DEFAULT_TOP_P = 1
+DEFAULT_PRESENCE_PENALTY = 0
+DEFAULT_FREQUENCY_PENALTY = 0
+DEFAULT_REPLY_COUNT = 1

embedding_model.py

@@ -0,0 +1,47 @@
+
+from base_class import Embedding_Model
+import pickle
+from sentence_transformers import SentenceTransformer
+
+from openai.embeddings_utils import (
+    get_embedding,
+)
+
+
+class HuggingfaceSentenceTransformerModel(Embedding_Model):
+    EMBEDDING_MODEL = "distiluse-base-multilingual-cased-v2"
+
+    def __init__(self, model_name=EMBEDDING_MODEL) -> None:
+        super().__init__(model_name)
+        
+        self.model = SentenceTransformer(model_name)
+
+    def __call__(self, text) -> None:
+        return self.model.encode(text)
+
+
+class OpenAIEmbeddingModel(Embedding_Model):
+    # constants
+    EMBEDDING_MODEL = "text-embedding-ada-002"
+    # establish a cache of embeddings to avoid recomputing
+    # cache is a dict of tuples (text, model) -> embedding, saved as a pickle file
+
+    def __init__(self, model_name=EMBEDDING_MODEL) -> None:
+        super().__init__(model_name)
+        self.model_name = model_name
+
+    # define a function to retrieve embeddings from the cache if present, and otherwise request via the API
+    def embedding_from_string(self,
+                              string: str,
+                              ) -> list:
+        """Return embedding of given string, using a cache to avoid recomputing."""
+        model = self.model_name
+        if (string, model) not in self.embedding_cache.keys():
+            self.embedding_cache[(string, model)] = get_embedding(
+                string, model)
+            with open(self.embedding_cache_path, "wb") as embedding_cache_file:
+                pickle.dump(self.embedding_cache, embedding_cache_file)
+        return self.embedding_cache[(string, model)]
+
+    def __call__(self, text) -> None:
+        return self.embedding_from_string(text)

frontend.py

@@ -0,0 +1,87 @@
+import datetime
+import os
+import streamlit as st
+from streamlit_chat import message
+import requests
+from config import PDF_SAVE_DIR
+
+st.set_page_config(
+    page_title="ChatPaper - Demo",
+    page_icon=":robot:"
+)
+
+pdf_uploaded = False
+
+if pdf_uploaded is False:
+    st.sidebar.markdown("## Upload a PDF")
+    pdf_uploader = st.sidebar.file_uploader("Upload a PDF", type="pdf", )
+
+st.sidebar.markdown("## API Key")
+api_key = st.sidebar.text_input(
+    "OpenAI API Key", value="", label_visibility="hidden", help="Please enter your API key.")
+
+
+def get_text():
+    input_text = st.text_input(
+        "User: ", "", help="Please ask any questions about the paper.")
+    return input_text
+
+
+st.header("ChatPaper - Demo")
+
+API_URL = "http://localhost:5000/query/"
+header = {"api_key": ""}
+
+if 'generated' not in st.session_state:
+    st.session_state['generated'] = []
+
+if 'past' not in st.session_state:
+    st.session_state['past'] = []
+
+if "user_stamp" not in st.session_state:
+    import datetime
+    user_stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    st.session_state['user_stamp'] = user_stamp
+
+
+if pdf_uploader is not None:
+    if api_key:
+        header['api_key'] = api_key
+        pdf_name = pdf_uploader.name.replace(' ', '_')
+
+        file_name = f"{st.session_state.user_stamp}_{pdf_name}"
+
+        # check PDF_SAVE_DIR
+        if not os.path.exists(PDF_SAVE_DIR):
+            os.makedirs(PDF_SAVE_DIR)
+
+        filepath = os.path.join(PDF_SAVE_DIR, file_name)
+        with open(filepath, "wb") as f:
+            f.write(pdf_uploader.getbuffer())
+        user_query = get_text()
+
+        if user_query:
+            st.session_state.past.append(user_query)
+            query_data = {"pdf_link": filepath,
+                          "user_stamp": st.session_state.user_stamp, "user_query": user_query}
+            print(query_data)
+            response = requests.post(
+                API_URL, headers=header, json=query_data, timeout=300)
+            output = response.json()
+            code = output['code']
+            response = output['response']
+            if code == 200:
+                st.session_state.generated.append(response)
+
+        if st.session_state['generated']:
+            for i in range(len(st.session_state['generated'])-1, -1, -1):
+                message(st.session_state["generated"][i],
+                        key=str(i), avatar_style="fun-emoji")
+                message(st.session_state['past'][i], is_user=True, key=str(
+                    i) + '_user', avatar_style="personas")
+    else:
+        st.markdown(
+            "<span style='color:red'>Please enter your API key.</span>", unsafe_allow_html=True)
+else:
+    st.markdown("<span style='color:red'>Please upload a PDF file.</span>",
+                unsafe_allow_html=True)

pdf_parser.py

@@ -0,0 +1,148 @@
+from base_class import AbstractPDFParser
+import pickle
+from scipdf_utils import parse_pdf_to_dict
+
+
+class GrobidSciPDFPaser(AbstractPDFParser):
+    import pysbd
+    seg_en = pysbd.Segmenter(language="en", clean=False)
+    seg_chinese = pysbd.Segmenter(language="zh", clean=False)
+
+    def __init__(self, pdf_link, db_name="grobid_scipdf", short_thereshold=30) -> None:
+        """Initialize the PDF parser
+
+            Args:
+                pdf_link: link to the PDF file, the pdf link can be a web link or local file path
+                metadata: metadata of the PDF file, like authors, title, abstract, etc.
+                paragraphs: list of paragraphs of the PDF file, all paragraphs are concatenated together
+                split_paragraphs: dict of section name and corresponding list of split paragraphs
+        """
+        super().__init__(db_name=db_name)
+        self.db_name = db_name
+        self.pdf_link = pdf_link
+        self.pdf = None
+        self.metadata = {}
+        self.flattn_paragraphs = None
+        self.split_paragraphs = None
+        self.short_thereshold = short_thereshold
+        self.parse_pdf()
+
+    def _contact_too_short_paragraphs(self, ):
+        """Contact too short paragraphs or discard them"""
+        for i, section in enumerate(self.split_paragraphs):
+            # section_name = section['heading']
+            paragraphs = section['texts']
+            new_paragraphs = []
+            for paragraph in paragraphs:
+                if len(paragraph) <= self.short_thereshold and len(paragraph.strip()) != 0:
+                    if len(new_paragraphs) != 0:
+                        new_paragraphs[-1] += paragraph
+                    else:
+                        new_paragraphs.append(paragraph)
+                else:
+                    new_paragraphs.append(paragraph)
+            self.split_paragraphs[i]['texts'] = new_paragraphs
+
+    @staticmethod
+    def _find_largest_font_string(file_name, search_string):
+        search_string = search_string.strip()
+        max_font_size = -1
+        page_number = -1
+        import PyPDF2
+        from pdfminer.high_level import extract_pages
+        from pdfminer.layout import LTTextContainer, LTChar
+        try: 
+            with open(file_name, 'rb') as file:
+                pdf_reader = PyPDF2.PdfReader(file)
+
+                for index, page_layout in enumerate(extract_pages(file_name)):
+                    for element in page_layout:
+                        if isinstance(element, LTTextContainer):
+                            for text_line in element:
+                                if search_string in text_line.get_text():
+                                    for character in text_line:
+                                        if isinstance(character, LTChar):
+                                            if character.size > max_font_size:
+                                                max_font_size = character.size
+                                                page_number = index
+            return page_number + 1 if page_number != -1 else -1
+        except Exception as e:
+            return -1
+        
+
+    def _find_section_page(self, section_name) -> None:
+        return GrobidSciPDFPaser._find_largest_font_string(self.pdf_link, section_name)
+
+    def _retrive_or_parse(self, ):
+        """Return pdf dict from cache if present, otherwise parse the pdf"""
+        db_name = self.db_name
+        if (self.pdf_link, db_name) not in self.db_cache.keys():
+            self.db_cache[(self.pdf_link, db_name)
+                          ] = parse_pdf_to_dict(self.pdf_link)
+            with open(self.db_cache_path, "wb") as db_cache_file:
+                pickle.dump(self.db_cache, db_cache_file)
+        return self.db_cache[(self.pdf_link, db_name)]
+
+    @staticmethod
+    def _check_chinese(text) -> None:
+        return any(u'\u4e00' <= char <= u'\u9fff' for char in text)
+
+    def parse_pdf(self) -> None:
+        """Parse the PDF file
+        """
+        article_dict = self._retrive_or_parse()
+        self.article_dict = article_dict
+        self._get_metadata()
+        self.split_paragraphs = self.get_split_paragraphs()
+        self._contact_too_short_paragraphs()
+
+        self.flattn_paragraphs = self.get_paragraphs()
+
+    def get_paragraphs(self) -> None:
+        """Get the paragraphs of the PDF file
+        """
+        paragraphs = []
+        self.content2section = {}
+        for section in self.split_paragraphs:
+            # paragraphs+=[section["heading"]]
+            paragraphs += section["texts"]
+            for para in section["texts"]:
+                self.content2section[para] = section["heading"]
+        return paragraphs
+
+    def _get_metadata(self) -> None:
+        for meta in ['authors', "pub_date", "abstract", "references", "doi", 'title',]:
+            self.metadata[meta] = self.article_dict[meta]
+        self.section_names = [section["heading"]
+                              for section in self.article_dict['sections']]
+        self.section_names2page = {}
+        for section_name in self.section_names:
+            section_page_index = self._find_section_page(section_name)
+            self.section_names2page.update({section_name: section_page_index})
+        self.section_names_with_page_index = [section_name + " (Page {})".format(
+            self.section_names2page[section_name]) for section_name in self.section_names]
+
+    def get_split_paragraphs(self, ) -> None:
+        section_pair_list = []
+        for section in self.article_dict['sections']:
+            section_pair_list.append({
+                "heading": section["heading"],
+                "texts": section["all_paragraphs"],
+            })
+        return section_pair_list
+
+    @staticmethod
+    def _determine_optimal_split_of_pargraphs(section_pair_list) -> None:
+        """
+        split based on the some magic rules
+        """
+        import pysbd
+        for section_pair in section_pair_list:
+            if GrobidSciPDFPaser._check_chinese(section_pair["text"]):
+                seg = GrobidSciPDFPaser.seg_chinese
+            else:
+                seg = GrobidSciPDFPaser.seg_en
+            section_pair["texts"] = seg.segment(section_pair["texts"])
+            section_pair["texts"] = [
+                para for para in section_pair["text"] if len(para) > 2]
+        return section_pair_list

requirements.txt

@@ -1 +1,7 @@
-flask
+Flask==2.2.3
+streamlit==1.20.0
+streamlit-chat==0.0.2.2
+torch
+sentence_transformers
+bs4
+openai

scipdf_utils.py

@@ -0,0 +1,424 @@
+import re
+import os
+import os.path as op
+from glob import glob
+import urllib
+import subprocess
+import requests
+from bs4 import BeautifulSoup, NavigableString
+
+
+# or https://cloud.science-miner.com/grobid/ for cloud service
+GROBID_URL = "http://localhost:8070"
+DIR_PATH = op.dirname(op.abspath(__file__))
+PDF_FIGURES_JAR_PATH = op.join(
+    DIR_PATH, "pdffigures2", "pdffigures2-assembly-0.0.12-SNAPSHOT.jar"
+)
+
+
+def list_pdf_paths(pdf_folder: str):
+    """
+    list of pdf paths in pdf folder
+    """
+    return glob(op.join(pdf_folder, "*", "*", "*.pdf"))
+
+
+def validate_url(path: str):
+    """
+    Validate a given ``path`` if it is URL or not
+    """
+    regex = re.compile(
+        r"^(?:http|ftp)s?://"  # http:// or https://
+        # domain...
+        r"(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|[A-Z0-9-]{2,}\.?)|"
+        r"localhost|"  # localhost...
+        r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})"  # ...or ip
+        r"(?::\d+)?"  # optional port
+        r"(?:/?|[/?]\S+)$",
+        re.IGNORECASE,
+    )
+    return re.match(regex, path) is not None
+
+
+def parse_pdf(
+    pdf_path: str,
+    fulltext: bool = True,
+    soup: bool = False,
+    grobid_url: str = GROBID_URL,
+):
+    """
+    Function to parse PDF to XML or BeautifulSoup using GROBID tool
+
+    You can see http://grobid.readthedocs.io/en/latest/Install-Grobid/ on how to run GROBID locally
+    After loading GROBID zip file, you can run GROBID by using the following
+    >> ./gradlew run
+
+    Parameters
+    ==========
+    pdf_path: str or bytes, path or URL to publication or article or bytes string of PDF
+    fulltext: bool, option for parsing, if True, parse full text of the article
+        if False, parse only header
+    grobid_url: str, url to GROBID parser, default at 'http://localhost:8070'
+        This could be changed to "https://cloud.science-miner.com/grobid/" for the cloud service
+    soup: bool, if True, return BeautifulSoup of the article
+
+    Output
+    ======
+    parsed_article: if soup is False, return parsed XML in text format,
+        else return BeautifulSoup of the XML
+    Example
+    =======
+    >> parsed_article = parse_pdf(pdf_path, fulltext=True, soup=True)
+    """
+    # GROBID URL
+    if fulltext:
+        url = "%s/api/processFulltextDocument" % grobid_url
+    else:
+        url = "%s/api/processHeaderDocument" % grobid_url
+
+    if isinstance(pdf_path, str):
+        if validate_url(pdf_path) and op.splitext(pdf_path)[-1].lower() != ".pdf":
+            print("The input URL has to end with ``.pdf``")
+            parsed_article = None
+        elif validate_url(pdf_path) and op.splitext(pdf_path)[-1] == ".pdf":
+            page = urllib.request.urlopen(pdf_path).read()
+            parsed_article = requests.post(url, files={"input": page}).text
+        elif op.exists(pdf_path):
+            parsed_article = requests.post(
+                url, files={"input": open(pdf_path, "rb")}
+            ).text
+        else:
+            parsed_article = None
+    elif isinstance(pdf_path, bytes):
+        # assume that incoming is byte string
+        parsed_article = requests.post(url, files={"input": pdf_path}).text
+    else:
+        parsed_article = None
+
+    if soup and parsed_article is not None:
+        parsed_article = BeautifulSoup(parsed_article, "lxml")
+    return parsed_article
+
+
+def parse_authors(article):
+    """
+    Parse authors from a given BeautifulSoup of an article
+    """
+    author_names = article.find("sourcedesc").findAll("persname")
+    authors = []
+    for author in author_names:
+        firstname = author.find("forename", {"type": "first"})
+        firstname = firstname.text.strip() if firstname is not None else ""
+        middlename = author.find("forename", {"type": "middle"})
+        middlename = middlename.text.strip() if middlename is not None else ""
+        lastname = author.find("surname")
+        lastname = lastname.text.strip() if lastname is not None else ""
+        if middlename != "":
+            authors.append(firstname + " " + middlename + " " + lastname)
+        else:
+            authors.append(firstname + " " + lastname)
+    authors = "; ".join(authors)
+    return authors
+
+
+def parse_date(article):
+    """
+    Parse date from a given BeautifulSoup of an article
+    """
+    pub_date = article.find("publicationstmt")
+    year = pub_date.find("date")
+    year = year.attrs.get("when") if year is not None else ""
+    return year
+
+
+def parse_abstract(article):
+    """
+    Parse abstract from a given BeautifulSoup of an article
+    """
+    div = article.find("abstract")
+    abstract = ""
+    for p in list(div.children):
+        if not isinstance(p, NavigableString) and len(list(p)) > 0:
+            abstract += " ".join(
+                [elem.text for elem in p if not isinstance(
+                    elem, NavigableString)]
+            )
+    return abstract
+
+
+def calculate_number_of_references(div):
+    """
+    For a given section, calculate number of references made in the section
+    """
+    n_publication_ref = len(
+        [ref for ref in div.find_all("ref") if ref.attrs.get("type") == "bibr"]
+    )
+    n_figure_ref = len(
+        [ref for ref in div.find_all(
+            "ref") if ref.attrs.get("type") == "figure"]
+    )
+    return {"n_publication_ref": n_publication_ref, "n_figure_ref": n_figure_ref}
+
+
+def parse_sections(article, as_list: bool = False):
+    """
+    Parse list of sections from a given BeautifulSoup of an article
+
+    Parameters
+    ==========
+    as_list: bool, if True, output text as a list of paragraph instead
+        of joining it together as one single text
+    """
+    article_text = article.find("text")
+    divs = article_text.find_all(
+        "div", attrs={"xmlns": "http://www.tei-c.org/ns/1.0"})
+    sections = []
+    for div in divs:
+        div_list = list(div.children)
+        if len(div_list) == 0:
+            heading = ""
+            text = ""
+            all_paragraphs = []
+        elif len(div_list) == 1:
+            if isinstance(div_list[0], NavigableString):
+                heading = str(div_list[0])
+                text = ""
+                all_paragraphs = []
+            else:
+                heading = ""
+                text = div_list[0].text
+                all_paragraphs = [text]
+        else:
+            text = []
+            heading = div_list[0]
+            all_paragraphs = []
+            if isinstance(heading, NavigableString):
+                heading = str(heading)
+                p_all = list(div.children)[1:]
+            else:
+                heading = ""
+                p_all = list(div.children)
+            for p in p_all:
+                if p is not None:
+                    try:
+                        text.append(p.text)
+                        all_paragraphs.append(p.text)
+                    except:
+                        pass
+            if not as_list:
+                text = "\n".join(text)
+        if heading != "" or text != "":
+            ref_dict = calculate_number_of_references(div)
+            sections.append(
+                {
+                    "heading": heading,
+                    "text": text,
+                    "all_paragraphs": all_paragraphs,
+                    "n_publication_ref": ref_dict["n_publication_ref"],
+                    "n_figure_ref": ref_dict["n_figure_ref"],
+                }
+            )
+    return sections
+
+
+def parse_references(article):
+    """
+    Parse list of references from a given BeautifulSoup of an article
+    """
+    reference_list = []
+    references = article.find("text").find("div", attrs={"type": "references"})
+    references = references.find_all(
+        "biblstruct") if references is not None else []
+    reference_list = []
+    for reference in references:
+        title = reference.find("title", attrs={"level": "a"})
+        if title is None:
+            title = reference.find("title", attrs={"level": "m"})
+        title = title.text if title is not None else ""
+        journal = reference.find("title", attrs={"level": "j"})
+        journal = journal.text if journal is not None else ""
+        if journal == "":
+            journal = reference.find("publisher")
+            journal = journal.text if journal is not None else ""
+        year = reference.find("date")
+        year = year.attrs.get("when") if year is not None else ""
+        authors = []
+        for author in reference.find_all("author"):
+            firstname = author.find("forename", {"type": "first"})
+            firstname = firstname.text.strip() if firstname is not None else ""
+            middlename = author.find("forename", {"type": "middle"})
+            middlename = middlename.text.strip() if middlename is not None else ""
+            lastname = author.find("surname")
+            lastname = lastname.text.strip() if lastname is not None else ""
+            if middlename != "":
+                authors.append(firstname + " " + middlename + " " + lastname)
+            else:
+                authors.append(firstname + " " + lastname)
+        authors = "; ".join(authors)
+        reference_list.append(
+            {"title": title, "journal": journal, "year": year, "authors": authors}
+        )
+    return reference_list
+
+
+def parse_figure_caption(article):
+    """
+    Parse list of figures/tables from a given BeautifulSoup of an article
+    """
+    figures_list = []
+    figures = article.find_all("figure")
+    for figure in figures:
+        figure_type = figure.attrs.get("type") or ""
+        figure_id = figure.attrs["xml:id"] or ""
+        label = figure.find("label").text
+        if figure_type == "table":
+            caption = figure.find("figdesc").text
+            data = figure.table.text
+        else:
+            caption = figure.text
+            data = ""
+        figures_list.append(
+            {
+                "figure_label": label,
+                "figure_type": figure_type,
+                "figure_id": figure_id,
+                "figure_caption": caption,
+                "figure_data": data,
+            }
+        )
+    return figures_list
+
+
+def convert_article_soup_to_dict(article, as_list: bool = False):
+    """
+    Function to convert BeautifulSoup to JSON format
+    similar to the output from https://github.com/allenai/science-parse/
+
+    Parameters
+    ==========
+    article: BeautifulSoup
+
+    Output
+    ======
+    article_json: dict, parsed dictionary of a given article in the following format
+        {
+            'title': ...,
+            'abstract': ...,
+            'sections': [
+                {'heading': ..., 'text': ...},
+                {'heading': ..., 'text': ...},
+                ...
+            ],
+            'references': [
+                {'title': ..., 'journal': ..., 'year': ..., 'authors': ...},
+                {'title': ..., 'journal': ..., 'year': ..., 'authors': ...},
+                ...
+            ],
+            'figures': [
+                {'figure_label': ..., 'figure_type': ..., 'figure_id': ..., 'figure_caption': ..., 'figure_data': ...},
+                ...
+            ]
+        }
+    """
+    article_dict = {}
+    if article is not None:
+        title = article.find("title", attrs={"type": "main"})
+        title = title.text.strip() if title is not None else ""
+        article_dict["authors"] = parse_authors(article)
+        article_dict["pub_date"] = parse_date(article)
+        article_dict["title"] = title
+        article_dict["abstract"] = parse_abstract(article)
+        article_dict["sections"] = parse_sections(article, as_list=as_list)
+        article_dict["references"] = parse_references(article)
+        article_dict["figures"] = parse_figure_caption(article)
+
+        doi = article.find("idno", attrs={"type": "DOI"})
+        doi = doi.text if doi is not None else ""
+        article_dict["doi"] = doi
+
+        return article_dict
+    else:
+        return None
+
+
+def parse_pdf_to_dict(
+    pdf_path: str,
+    fulltext: bool = True,
+    soup: bool = True,
+    as_list: bool = False,
+    grobid_url: str = GROBID_URL,
+):
+    """
+    Parse the given PDF and return dictionary of the parsed article
+
+    Parameters
+    ==========
+    pdf_path: str, path to publication or article
+    fulltext: bool, whether to extract fulltext or not
+    soup: bool, whether to return BeautifulSoup or not
+    as_list: bool, whether to return list of sections or not
+    grobid_url: str, url to grobid server, default is `GROBID_URL`
+        This could be changed to "https://cloud.science-miner.com/grobid/" for the cloud service
+
+    Ouput
+    =====
+    article_dict: dict, dictionary of an article
+    """
+    parsed_article = parse_pdf(
+        pdf_path, fulltext=fulltext, soup=soup, grobid_url=grobid_url
+    )
+    article_dict = convert_article_soup_to_dict(
+        parsed_article, as_list=as_list)
+    return article_dict
+
+
+def parse_figures(
+    pdf_folder: str,
+    jar_path: str = PDF_FIGURES_JAR_PATH,
+    resolution: int = 300,
+    output_folder: str = "figures",
+):
+    """
+    Parse figures from the given scientific PDF using pdffigures2
+
+    Parameters
+    ==========
+    pdf_folder: str, path to a folder that contains PDF files. A folder must contains only PDF files
+    jar_path: str, default path to pdffigures2-assembly-0.0.12-SNAPSHOT.jar file
+    resolution: int, resolution of the output figures
+    output_folder: str, path to folder that we want to save parsed data (related to figures) and figures
+
+    Output
+    ======
+    folder: making a folder of output_folder/data and output_folder/figures of parsed data and figures relatively
+    """
+    if not op.isdir(output_folder):
+        os.makedirs(output_folder)
+
+    # create ``data`` and ``figures`` subfolder within ``output_folder``
+    data_path = op.join(output_folder, "data")
+    figure_path = op.join(output_folder, "figures")
+    if not op.exists(data_path):
+        os.makedirs(data_path)
+    if not op.exists(figure_path):
+        os.makedirs(figure_path)
+
+    if op.isdir(data_path) and op.isdir(figure_path):
+        args = [
+            "java",
+            "-jar",
+            jar_path,
+            pdf_folder,
+            "-i",
+            str(resolution),
+            "-d",
+            os.path.join(os.path.abspath(data_path), ""),
+            "-m",
+            op.join(os.path.abspath(figure_path), ""),  # end path with "/"
+        ]
+        _ = subprocess.run(
+            args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=20
+        )
+        print("Done parsing figures from PDFs!")
+    else:
+        print("You may have to check of ``data`` and ``figures`` in the the output folder path.")

serve_grobid.sh

@@ -11,4 +11,4 @@ fi
 
 # run GROBID
 cd grobid-${GROBID_VERSION} || exit
-./gradlew run &
+./gradlew run

similarity_metric.py

@@ -0,0 +1,13 @@
+from base_class import SimilarityAlg
+from openai.embeddings_utils import (
+    distances_from_embeddings,
+)
+
+
+class CosineSimilarity(SimilarityAlg):
+    def __init__(self) -> None:
+        pass
+
+    @staticmethod
+    def __call__(query_embedding, embeddings) -> None:
+        return distances_from_embeddings(query_embedding, embeddings, distance_metric="cosine")

utils.py

@@ -0,0 +1,24 @@
+
+from typing import Set
+
+
+def get_filtered_keys_from_object(obj: object, *keys: str) -> Set[str]:
+    """
+    Get filtered list of object variable names.
+    :param keys: List of keys to include. If the first key is "not", the remaining keys will be removed from the class keys.
+    :return: List of class keys.
+    """
+    class_keys = obj.__dict__.keys()
+    if not keys:
+        return class_keys
+
+    # Remove the passed keys from the class keys.
+    if keys[0] == "not":
+        return {key for key in class_keys if key not in keys[1:]}
+    # Check if all passed keys are valid
+    if invalid_keys := set(keys) - class_keys:
+        raise ValueError(
+            f"Invalid keys: {invalid_keys}",
+        )
+    # Only return specified keys that are in class_keys
+    return {key for key in keys if key in class_keys}