Legora / db /all_chunks.py
sai-Rohan's picture
combined results of all parsers and made a single pipeline
dbabef2
Raw
History Blame Contribute Delete
6.21 kB
from __future__ import annotations
# =====================================================
# BNS
# =====================================================
from db.parsers.bns.cleaner import (
BNSTextCleaner
)
from db.parsers.bns.bns_parser import (
BNSParser
)
from db.parsers.bns.chunker import (
LegalChunker as BNSChunker,
chunks_to_dicts as bns_to_dicts
)
# =====================================================
# BNSS
# =====================================================
from db.parsers.bnss.cleaner import (
BNSSTextCleaner
)
from db.parsers.bnss.bnss_parser import (
BNSSParser
)
from db.parsers.bnss.chunker import (
LegalChunker as BNSSChunker,
chunks_to_dicts as bnss_to_dicts
)
# =====================================================
# BSA
# =====================================================
from db.parsers.bsa.cleaner import (
BSATextCleaner
)
from db.parsers.bsa.bsa_parser import (
BSAParser
)
from db.parsers.bsa.chunker import (
LegalChunker as BSAChunker,
chunks_to_dicts as bsa_to_dicts
)
# =====================================================
# CONSTITUTION
# =====================================================
from db.parsers.consitution.cleaner import (
ConstitutionTextCleaner
)
from db.parsers.consitution.constitution_parser import (
ConstitutionParser
)
from db.parsers.consitution.chunker import (
LegalChunker as ConstitutionChunker,
chunks_to_dicts as constitution_to_dicts
)
class LegalCorpusBuilder:
# =================================================
# LOAD FILE
# =================================================
@staticmethod
def load_text(
filepath
):
with open(
filepath,
"r",
encoding="utf8"
) as f:
return f.read()
# =================================================
# BNS
# =================================================
def build_bns_chunks(
self,
filepath
):
text = self.load_text(
filepath
)
text = (
BNSTextCleaner()
.clean(text)
)
document = (
BNSParser()
.parse(text)
)
chunks = (
BNSChunker(
document_name="BNS"
)
.chunk_document(
document
)
)
return (
bns_to_dicts(
chunks
)
)
# =================================================
# BNSS
# =================================================
def build_bnss_chunks(
self,
filepath
):
text = self.load_text(
filepath
)
text = (
BNSSTextCleaner()
.clean(text)
)
document = (
BNSSParser()
.parse(text)
)
chunks = (
BNSSChunker(
document_name="BNSS"
)
.chunk_document(
document
)
)
return (
bnss_to_dicts(
chunks
)
)
# =================================================
# BSA
# =================================================
def build_bsa_chunks(
self,
filepath
):
text = self.load_text(
filepath
)
text = (
BSATextCleaner()
.clean(text)
)
document = (
BSAParser()
.parse(text)
)
chunks = (
BSAChunker(
document_name="BSA"
)
.chunk_document(
document
)
)
return (
bsa_to_dicts(
chunks
)
)
# =================================================
# CONSTITUTION
# =================================================
def build_constitution_chunks(
self,
filepath
):
text = self.load_text(
filepath
)
text = (
ConstitutionTextCleaner()
.clean(text)
)
document = (
ConstitutionParser()
.parse(text)
)
chunker = ConstitutionChunker()
chunks = (
chunker.chunk_constitution(
document
)
)
return (
constitution_to_dicts(
chunks
)
)
# =================================================
# ALL CHUNKS
# =================================================
def build_all_chunks(
self,
bns_path,
bnss_path,
bsa_path,
constitution_path
):
all_chunks = []
print(
"\nBuilding BNS..."
)
all_chunks.extend(
self.build_bns_chunks(
bns_path
)
)
print(
"Building BNSS..."
)
all_chunks.extend(
self.build_bnss_chunks(
bnss_path
)
)
print(
"Building BSA..."
)
all_chunks.extend(
self.build_bsa_chunks(
bsa_path
)
)
print(
"Building Constitution..."
)
all_chunks.extend(
self.build_constitution_chunks(
constitution_path
)
)
return all_chunks
# =====================================================
# TEST
# =====================================================
if __name__ == "__main__":
builder = (
LegalCorpusBuilder()
)
all_chunks = (
builder.build_all_chunks(
bns_path=
"db/pdfs/bns.txt",
bnss_path=
"db/pdfs/bnss.txt",
bsa_path=
"db/pdfs/bsa.txt",
constitution_path=
"db/pdfs/constitution.txt"
)
)
print()
print(
f"Total Chunks: "
f"{len(all_chunks)}"
)