Upload 7 files

.gitignore

@@ -0,0 +1,3 @@
+flagged/
+.idea/
+*/__pycache__/

main.py

@@ -0,0 +1,38 @@
+# This is a sample Python script.
+
+# Press Shift+F10 to execute it or replace it with your code.
+# Press Double Shift to search everywhere for classes, files, tool windows, actions, and settings.
+
+import gradio as gr
+import os
+
+from reranker.reranker import CrossEncReranker
+from retriever.es_retriever import ESRetriever
+from utils.preprocessing import question_to_statement
+
+
+ES_HOST = os.environ["ES_HOST"]
+ES_INDEX_NAME = os.environ["ES_INDEX_NAME"]
+ES_USERNAME = os.environ["ES_USERNAME"]
+ES_PASSWORD = os.environ["ES_PASSWORD"]
+
+RERANKER_MODEL_NAME = "douglasfaisal/granularity-legal-reranker-cross-encoder-indobert-base-p2"
+
+es_retriever_client = ESRetriever(ES_HOST, ES_INDEX_NAME, ES_USERNAME, ES_PASSWORD)
+cross_enc_reranker = CrossEncReranker(RERANKER_MODEL_NAME, 512)
+
+def retrieve_and_rerank(question: str):
+
+    query = question_to_statement(question)
+    retrieval_results = es_retriever_client.retrieve(query)
+    reranker_results = cross_enc_reranker.rerank(query, retrieval_results)
+
+    return reranker_results[0].text
+
+
+demo = gr.Interface(fn=retrieve_and_rerank, inputs="text", outputs="text")
+
+# Press the green button in the gutter to run the script.
+demo.launch()
+
+# See PyCharm help at https://www.jetbrains.com/help/pycharm/

models/law_component.py

@@ -0,0 +1,153 @@
+import copy
+
+LAW_PATTERN = '(?P<law_type>[A-Za-z]+) (?P<law_number>\d{1,})/(?P<law_year>\d{4})'
+
+class LawComponent:
+    law_type: str = None
+    law_year: int = 0
+    law_number: int = 0
+    component_type: str = None
+    chapter: int = None
+    article: str = None
+    subsection: int = None
+    letter: str = None
+    text: str = None
+
+    def __init__(self):
+        pass
+
+    def from_uri(uri: str):
+        lc = LawComponent()
+        uri_split = uri.split('/')
+        lc.law_type = uri_split[4]
+        lc.law_year = int(uri_split[5]) if int(uri_split[5]) != 0 else None
+        lc.law_number = int(uri_split[6]) if int(uri_split[6]) != 0 else None
+        if len(uri_split) < 8:
+            return lc
+        if (uri_split[7] == 'bab'):
+            lc.component_type = 'chapter'
+            lc.chapter = int(uri_split[8]) if int(uri_split[8]) != 0 else None
+        else:
+            lc.article = str(int(uri_split[8])) if int(uri_split[8]) != 0 else None
+            lc.component_type = 'article'
+            if (len(uri_split) > 9 and uri_split[9] == "versi"):
+                if (len(uri_split) > 11 and uri_split[11] == "ayat"):
+                    lc.subsection = int(uri_split[12]) if int(uri_split[12]) != 0 else None
+                    lc.component_type = 'subsection'
+                    if (len(uri_split) > 13 and uri_split[13] == "huruf"):
+                        lc.component_type = 'letter'
+                        try:
+                            lc.letter = str(int(uri_split[14])) if int(uri_split[14]) != 0 else None
+                        except:
+                            lc.letter = uri_split[14]
+                elif (len(uri_split) > 11 and uri_split[11] == "huruf"):
+                    lc.component_type = 'letter'
+                    try:
+                        lc.letter = str(int(uri_split[12])) if int(uri_split[12]) != 0 else None
+                    except:
+                        lc.letter = uri_split[12]
+
+        return lc
+
+    # def from_answer_granularity_row(row_dict: dict):
+    #     lc = LawComponent()
+    #
+    #     law = row_dict['Law']
+    #     law_search = re.search(LAW_PATTERN, law)
+    #
+    #     if (law_search != None):
+    #         lc.law_type = law_search.group('law_type').lower()
+    #         lc.law_number = int(law_search.group('law_number')) if int(law_search.group('law_number')) != 0 else None
+    #         lc.law_year = int(law_search.group('law_year')) if int(law_search.group('law_year')) != 0 else None
+    #
+    #         lc.component_type = row_dict['Answer Granularity'].lower()
+    #         lc.chapter = int(row_dict['Chapter']) if int(row_dict['Chapter']) != 0 else None
+    #         try:
+    #             lc.article = str(int(row_dict['Article'])) if int(row_dict['Article']) != 0 else None
+    #         except:
+    #             pass
+    #         try:
+    #             lc.subsection = int(row_dict['Subsection']) if int(row_dict['Subsection']) != 0 else None
+    #         except:
+    #             pass
+    #         try:
+    #             lc.letter = str(int(row_dict['Letter (1st level)'])) if int(
+    #                 row_dict['Letter (1st level)']) != 0 else None
+    #         except:
+    #             if (pd.isnull(row_dict['Letter (1st level)'])):
+    #                 lc.letter = None
+    #             else:
+    #                 lc.letter = row_dict['Letter (1st level)']
+    #
+    #     return lc
+
+    def set_text(self, text):
+        self.text = text
+
+    def __eq__(self, other):
+
+        if not (self.law_type == other.law_type and self.law_year == other.law_year
+                and self.law_number == other.law_number):
+            return False
+
+        if self.component_type != other.component_type:
+            return False
+
+        # if self.component_type == 'chapter':
+        #   if self.chapter != other.chapter:
+        #     return False
+        if self.article is None and other.article is None:
+            if self.chapter != other.chapter:
+                return False
+        else:
+            if self.article != other.article:
+                return False
+
+            if self.component_type == 'article':
+                return True
+
+            if self.subsection != other.subsection:
+                return False
+
+            if self.component_type == 'subsection':
+                return True
+
+            if self.letter != other.letter:
+                return False
+
+            if self.component_type == 'letter':
+                return True
+
+        return True
+
+    def is_article_equal(self, other):
+
+        if not (self.law_type == other.law_type and self.law_year == other.law_year
+                and self.law_number == other.law_number):
+            return False
+
+        if self.component_type == 'chapter':
+            return False
+        else:
+            if self.article != other.article:
+                return False
+
+        return True
+
+    def __repr__(self):
+        return "LawComponent({}, {}, {}, {}, {}, {}, {}, {})".format(
+            self.law_type,
+            self.law_number,
+            self.law_year,
+            self.component_type,
+            self.chapter,
+            self.article,
+            self.subsection,
+            self.letter
+        )
+
+    def __str__(self):
+        return self.__repr__()
+
+    def copy(self):
+        return copy.deepcopy(self)

requirements.txt

@@ -0,0 +1,3 @@
+gradio~=3.28.1
+numpy==1.21.4
+requests==2.26.0

reranker/reranker.py

@@ -0,0 +1,20 @@
+import numpy as np
+
+from models.law_component import LawComponent
+from sentence_transformers.cross_encoder import CrossEncoder
+
+class CrossEncReranker:
+
+    def __init__(self, model_name, max_length=512):
+        self.model_name = model_name
+        self.reranker = CrossEncoder(self.model_name)
+        self.reranker.max_length = max_length
+
+    def rerank(self, query_text: str, candidates: list[LawComponent]):
+        sentence_combinations = [[query_text, c.text] for c in candidates]
+        similarity_scores = self.reranker.predict(sentence_combinations)
+        index = np.argsort(similarity_scores)[::-1]
+
+        reranked_candidates = np.array(candidates)[index]
+        return reranked_candidates
+

retriever/es_retriever.py

@@ -0,0 +1,89 @@
+import copy
+
+import json
+import requests
+from requests.auth import HTTPBasicAuth
+from models.law_component import LawComponent
+
+base_query = {
+    "query": {
+        "bool": {
+            "should": [
+                {
+                    "match": {
+                        "text": {
+                            "query": None,
+                            "boost": 1.0
+                        }
+                    }
+                },
+                {
+                    "match": {
+                        "chapterTitle": {
+                            "query": None,
+                            "boost": 1.0
+                        }
+                    }
+                },
+                {
+                    "match_phrase": {
+                        "text": {
+                            "query": None,
+                            "boost": 1.0
+                        }
+                    }
+                },
+                {
+                    "match_phrase": {
+                        "chapterTitle": {
+                            "query": None,
+                            "boost": 1.0
+                        }
+                    }
+                }
+            ],
+            "minimum_should_match": 1
+        }
+    }
+}
+
+class ESRetriever:
+
+    def __init__(self, es_host, es_index_name, es_username="", es_password=""):
+        self.es_host = es_host
+        self.es_index_name = es_index_name
+        self.es_username = es_username
+        self.es_password = es_password
+
+        if (es_username != "" and es_password != ""):
+            self.auth = HTTPBasicAuth(es_username, es_password)
+        else:
+            self.auth = None
+
+    # Returns LawComponent
+    def retrieve(self, query_text: str):
+        query = copy.deepcopy(base_query)
+        query['query']['bool']['should'][0]['match']['text']['query'] = query_text
+        query['query']['bool']['should'][1]['match']['chapterTitle']['query'] = query_text
+        query['query']['bool']['should'][2]['match_phrase']['text']['query'] = query_text
+        query['query']['bool']['should'][3]['match_phrase']['chapterTitle']['query'] = query_text
+
+        # try:
+        response = requests.get(
+            self.es_host + self.es_index_name + '/_search',
+            headers={'Content-Type': 'application/json'},
+            data=json.dumps(query),
+            auth=self.auth
+        )
+
+        if response.ok:
+            results = response.json()["hits"]["hits"]
+            retrieval_results = []
+            for result in results:
+                lc = LawComponent.from_uri(result["_source"]["uri"])
+                lc.set_text(result["_source"]["text"])
+                retrieval_results.append(lc)
+            return retrieval_results
+
+        #
+        # response.content

utils/preprocessing.py

@@ -0,0 +1,11 @@
+import re
+
+
+def question_to_statement(question: str):
+    pattern = "(?P<qWord>kapan|(ber|meng|si)?apa|(di|ke) ?(mana)|bagaimana) ?(saja)? ?(kah)?"
+
+    result = re.sub(pattern, "", question.lower())
+    result = result.replace("?", "")
+    result = re.sub("\s+", " ", result)
+    result = result.strip()
+    return result