Ask Shamela

ask_shamela.py

"""
env/bin/pip install huggingface-hub flashrank langchain-chroma langchain-ollama ollama langchain-openai dotenv

Contoh mengunduh kategori 24 (Biografi Nabi):

env/bin/hf download AuthenticIlm/Shamela4_Full_DB \
  --repo-type dataset \
  --include "24__*" \
  --local-dir maktabah_syamilah

Jadi kategori adalah batas pengetahuannya.

Bila menggunakan ChatGPT maka buat file .env berisi:

OPENAI_API_KEY=api-key-milikmu

Jalankan:

env/bin/python ask_shamela.py --category-id=24 --prompt="Siapa saja anak-anak Nabi ?"
"""
import sys
import os
import json
from time import time
from argparse import ArgumentParser
from glob import glob
import numpy as np
import chromadb
from langchain_core.documents import Document
from langchain_chroma import Chroma
from langchain_ollama import (
    OllamaEmbeddings,
    ChatOllama,
    )
from langchain_openai import (
    OpenAIEmbeddings,
    ChatOpenAI,
    )
from langchain_core.messages import (
    SystemMessage,
    HumanMessage,
    AIMessage,
    )
from flashrank import (
    Ranker,
    RerankRequest,
    )
import ollama
from dotenv import load_dotenv


# Sumber: https://huggingface.co/datasets/AuthenticIlm/Shamela4_Full_DB
shamela_dir = "maktabah_syamilah"
help_shamela_dir = f"default {shamela_dir}"

# Bila beberapa lama kemudian ada ResponseError maka itu
# artinya GPU kekurangan memori. Solusinya turunkan batch_size.
batch_size = 100
help_batch_size = (
    f"default {batch_size}, jumlah dokumen yang dikirim ke embed model")

related_doc_count = 25
help_related_doc_count = (
    f"default {related_doc_count}, jumlah dokumen terkait")

# Untuk LLM usai Rerank mengurutkannya
best_doc_count = 10
help_best_doc_count = (
    f"default {best_doc_count}, jumlah dokumen yang paling terkait")

# Command Line Argument
pars = ArgumentParser()
pars.add_argument("--shamela-dir", default=shamela_dir, help=help_shamela_dir)
pars.add_argument("--category-id", required=True, type=int)
pars.add_argument("--chroma-dir")
pars.add_argument("--prompt")
pars.add_argument(
    "--batch-size", type=int, default=batch_size, help=help_batch_size)
pars.add_argument(
    "--related-doc-count", type=int, default=related_doc_count,
    help=help_related_doc_count)
pars.add_argument(
    "--best-doc-count", type=int, default=best_doc_count,
    help=help_best_doc_count)
option = pars.parse_args(sys.argv[1:])

if option.chroma_dir:
    chroma_dir = option.chroma_dir
else:
    chroma_dir = f"chroma_{option.shamela_dir}"
collection_name = f"category_{option.category_id}"

RANKER_MODEL = "ms-marco-MultiBERT-L-12"
RANKER_DIR = "ranker"
SYSTEM_PROMPT_TEMPLATE = (
    "Jawabanmu harus semata-mata berasal dari basis pengetahuan yang "
    "disediakan di bawah ini.\n"
    "DIWAJIBKAN: Setiap kali kamu mengambil informasi dari sebuah "
    "paragraf, sebutkan dari BUKU ID mana informasi tersebut berasal "
    "di akhir kalimat atau paragraf jawabanmu! Cara menyebutkannya di "
    "dalam tanda kurung, contoh: (BUKU 1, HALAMAN 2)\n\n"
    "Basis Pengetahuan:\n{context}\n\n")
SYSTEM_PROMPT_DONT_KNOW = SystemMessage(
    "Jika informasi tidak tersedia maka jawaban harus berisi "
    "kalimat persis ini: \"Maaf, saya tidak mengetahui hal tersebut.\"")
FAILED_FILE = "book_ids.txt"
# Urutan sesuai https://shamela.ws/#categories
CATEGORIES = [
    "Doktrin",  # 01
    "Perbedaan dan Tanggapan",  # 02
    "Interpretasi",  # 03
    "Ilmu-ilmu Al-Quran dan Prinsip-prinsip Penafsiran",  # 04
    "Tajwid dan Bacaan Al-Quran",  # 05
    "Kitab-kitab Sunnah",  # 06
    "Penjelasan Hadits",  # 07
    "Kelulusan dan Pengesahan",  # 08
    "Kekurangan dan Pertanyaan tentang Hadits",  # 09
    "Ilmu Hadits",  # 10
    "Prinsip-prinsip Fiqih Islam",  # 11
    "Yurisprudensi dan Maksim Hukum",  # 12
    "Logika",  # 13
    "Fiqih Hanafi",  # 14
    "Yurisprudensi Maliki",  # 15
    "Fikih Syafi'i",  # 16
    "Yurisprudensi Hanbali",  # 17
    "Yurisprudensi Umum",  # 18
    "Isu-isu Yurisprudensi",  # 19
    "Teori Politik Islam dan Peradilan",  # 20
    "Kewajiban dan Wasiat",  # 21
    "Fatwa",  # 22
    "Penyempurnaan Rohani, Tata Krama, dan Pengingatan",  # 23
    "Biografi Nabi",  # 24
    "Sejarah",  # 25
    "Biografi dan Kelas",  # 26
    "Silsilah dan Keluarga",  # 27
    "Negara dan Perjalanan",  # 28
    "Buku-buku Bahasa",  # 29
    "Kata-kata Langka dan Kamus",  # 30
    "Tata Bahasa dan Morfologi",  # 31
    "Sastra",  # 32
    "Prosodi dan Sajak",  # 33
    "Puisi",  # 34
    "Retorika",  # 35
    "Masjid",  # 36
    "Indeks Buku dan Panduan",  # 37
    "Kedokteran",  # 38
    "Buku Umum",  # 39
    "Ilmu Pengetahuan Lainnya",  # 40
    ]


def get_category_dir() -> str:
    prefix = os.path.join(option.shamela_dir, f"{option.category_id}__")
    pattern = f"{prefix}*"
    dir_name = glob(pattern)
    if not dir_name:
        print(f"Tidak ada direktori yang berawalan {prefix}")
        sys.exit(1)
    return dir_name[0]


def patched_run(output_names, input_feed, run_options=None):
    # 1. Pastikan input_ids dan attention_mask berformat 2D
    for key in ['input_ids', 'attention_mask']:
        if key in input_feed and input_feed[key].ndim == 1:
            input_feed[key] = np.expand_dims(input_feed[key], axis=0)
    # 2. Suntikkan token_type_ids jika belum ada
    if 'input_ids' in input_feed and 'token_type_ids' not in input_feed:
        input_ids_shape = input_feed['input_ids'].shape
        input_feed['token_type_ids'] = np.zeros(
            input_ids_shape, dtype=np.int64)
    if 'token_type_ids' in input_feed and \
            input_feed['token_type_ids'].ndim == 1:
        input_feed['token_type_ids'] = np.expand_dims(
            input_feed['token_type_ids'], axis=0)
    # 3. ANTISIPASI TOKEN KOSONG (Panjang sekuens = 0)
    # Jika baris teks kosong/gagal ditokenisasi, bentuknya menjadi [1, 0]
    # Kita paksa ganti dengan token dummy [101, 102] (CLS & SEP standar
    # BERT) agar Gather node tidak error
    if 'input_ids' in input_feed and input_feed['input_ids'].shape[1] == 0:
        input_feed['input_ids'] = np.array([[101, 102]], dtype=np.int64)
        if 'attention_mask' in input_feed:
            input_feed['attention_mask'] = np.array([[1, 1]], dtype=np.int64)
        if 'token_type_ids' in input_feed:
            input_feed['token_type_ids'] = np.array([[0, 0]], dtype=np.int64)
    return ranker.session.run(output_names, input_feed, run_options)


def get_book_ids() -> list:
    book_ids = []
    for book_dir in os.listdir(category_dir):
        t = book_dir.split("__")
        book_id = t[0]
        book_id = int(book_id)
        book_ids.append(book_id)
    book_ids.sort()
    return book_ids


def create_data() -> dict:
    docs = []
    ids = []
    for book_id in book_ids:
        prefix = f"{book_id}__"
        matches = glob(os.path.join(category_dir, f"{prefix}*"))
        book_dir = matches[0]
        pages_file = os.path.join(book_dir, "pages.jsonl")
        print(f"Membaca file {pages_file} ...")
        with open(pages_file, encoding="utf-8") as f:
            for index, line in enumerate(f):
                data = json.loads(line)
                # Ekstrak data teks per halaman teks Shamela
                # Catatan: Sesuaikan key 'body' atau 'text' tergantung isi
                # jsonl asli repositori
                text = data.get("body", data.get("text", "")).strip()
                if not text:
                    continue  # Lewati jika halaman kosong
                page_num = data["page_num"]
                # Kumpulkan dokumen dan metadata penting untuk sitasi RAG nanti
                metadata = dict(book_id=book_id, page_num=page_num)
                doc = Document(page_content=text, metadata=metadata)
                docs.append(doc)
                line_id = index + 1
                doc_id = f"book_{book_id}_line_{line_id}"
                ids.append(doc_id)
    return dict(docs=docs, ids=ids)


def create_vector_store():
    print(
        f"Siapkan database vektor di direktori {chroma_dir} "
        f"koleksi {collection_name} ...")
    return Chroma(
        persist_directory=chroma_dir,
        collection_name=collection_name,
        embedding_function=embedding_client)


def get_book_count() -> int:
    book_ids = set()
    offset = 0
    limit = 2000
    while True:
        data = vector_store.get(
                include=["metadatas"], limit=limit, offset=offset)
        if not data["metadatas"]:
            break
        for m in data["metadatas"]:
            if m and "book_id" in m:
                book_ids.add(m["book_id"])
        offset += limit
    return len(book_ids)


def humanize_time(secs: int) -> str:
    mins, secs = divmod(secs, 60)
    hours, mins = divmod(mins, 60)
    return '%02d:%02d:%02d' % (hours, mins, secs)


def get_token(meta):
    if "token_usage" in meta:  # OpenAI invoke
        usage = meta["token_usage"]
        token_input = usage["prompt_tokens"]
        token_output = usage["completion_tokens"]
    elif "input_tokens" in meta:  # OpenAI stream
        token_input = meta["input_tokens"]
        token_output = meta["output_tokens"]
    else:  # Ollama
        token_input = meta["prompt_eval_count"]
        token_output = meta["eval_count"]
    return dict(input=token_input, output=token_output)


class BaseProgress:
    def __init__(self, data, batch_size=100):
        self.data = data
        self.batch_size = batch_size
        self.term_size = os.get_terminal_size().columns

    def run(self):
        self.count = len(self.data["docs"])
        self.count_done = 0
        self.est_seconds = None
        begin_time = time()
        for i in range(0, self.count, self.batch_size):
            batch = self.data["docs"][i:i + self.batch_size]
            batch_count = len(batch)
            self.count_done += batch_count
            self.process(batch, i)
            duration = time() - begin_time
            self.speed = duration / self.count_done
            remain = self.count - self.count_done
            self.est_seconds = remain * self.speed
        self.end()

    def process(self, batch: list, index: int):  # Override, please
        estimate = self.get_estimate_msg()
        msg = f"\rBatch {self.count_done}/{self.count}{estimate}"
        msg = self.msg_full_size(msg)
        print(msg, end="", flush=True)

    def get_estimate_msg(self):
        if self.est_seconds:
            seconds = int(self.est_seconds)
            estimate = humanize_time(self.est_seconds)
            return f" {self.speed:.2f} detik/doc, perkiraan {estimate} lagi"
        return ""

    def msg_full_size(self, s: str) -> str:
        spaces = " " * (self.term_size - len(s))
        return f"{s}{spaces}"

    def end(self):
        if self.est_seconds is not None:
            msg = (
                f"\rBatch {self.count}/{self.count}, "
                f"{self.speed:.2f} detik/doc, selesai.")
            msg = self.msg_full_size(msg)
            print(msg, end="", flush=True)
        print()


def book_ids_from_batch(batch: list) -> list:
    book_ids = set()
    for b in batch:
        book_ids.add(b.metadata["book_id"])
    return list(book_ids)


class ChromaProgress(BaseProgress):
    def process(self, batch: list, index: int):  # Override
        estimate = self.get_estimate_msg()
        book_ids = book_ids_from_batch(batch)
        msg = (
            f"\rBook {book_ids}, "
            f"Batch {self.count_done}/{self.count}{estimate}")
        msg = self.msg_full_size(msg)
        print(msg, end="", flush=True)
        book_ids = book_ids_from_batch(batch)
        d = [str(x) for x in book_ids]
        s = ",".join(d)
        with open(FAILED_FILE, "w") as f:
            f.write(s)
        ids = self.data["ids"][index:index + self.batch_size]
        vector_store.add_documents(documents=batch, ids=ids)
        os.remove(FAILED_FILE)


def delete_corrupted_books() -> list:
    with open(FAILED_FILE) as f:
        s = f.read()
    ids = []
    for book_id in s.split(","):
        book_id = int(book_id)
        print(f"Hapus buku {book_id} dari {chroma_dir} ...")
        vector_store.delete(where=dict(book_id=book_id))
        ids.append(book_id)
    return ids


def show_sample(docs):
    count = len(docs)
    print(
        f"Ada {count} potongan dokumen yang siap dimasukkan ke "
        f"direktori {chroma_dir}")
    print("\n--- Contoh Potongan 1 ---")
    doc = docs[0]
    print("Metadata")
    print(doc.metadata)
    print("Content")
    print(doc.page_content)
    print()


def ollama_embedding_class():
    url = os.getenv("URL", "http://localhost:11434")
    embedding_options["base_url"] = os.getenv("EMBEDDING_URL", url)
    embedding_options["model"] = os.getenv("EMBEDDING_MODEL", "bge-m3")
    return OllamaEmbeddings


def openai_embedding_class():
    embedding_options["model"] = os.getenv(
        "EMBEDDING_MODEL", "text-embedding-3-large")
    return OpenAIEmbeddings


def ollama_llm_class():
    url = os.getenv("URL", "http://localhost:11434")
    llm_options["base_url"] = os.getenv("LLM_URL", url)
    llm_options["model"] = os.getenv("LLM_MODEL", "gemma4:e4b")
    client = ollama.Client(host=llm_options["base_url"])
    info = client.show(llm_options["model"])
    for key, val in info.modelinfo.items():
        if key.find("context_length") > -1:
            # Maksimalkan kapasitas token
            llm_options["num_ctx"] = val
            print(f"dengan konteks {val} token")
            break
    return ChatOllama


def openai_llm_class():
    llm_options["model"] = os.getenv("LLM_MODEL", "gpt-5.1")
    llm_options["stream_usage"] = True
    return ChatOpenAI


# Load environment
load_dotenv()
embedding_options = dict()
llm_options = dict(temperature=0)
is_openai = os.getenv("OPENAI_API_KEY")
if is_openai:
    embedding_class = os.getenv("EMBEDDING_URL") and \
        ollama_embedding_class() or openai_embedding_class()
    llm_class = os.getenv("LLM_URL") and ollama_llm_class() or \
        openai_llm_class()
else:
    embedding_class = ollama_embedding_class()
    llm_class = ollama_llm_class()
print(f"Embedding options: {embedding_options}")
embedding_client = embedding_class(**embedding_options)
print(f"LLM options: {llm_options}")
llm_client = llm_class(**llm_options)

category_dir = get_category_dir()
book_ids = get_book_ids()
vector_store = create_vector_store()
client = chromadb.PersistentClient(path=chroma_dir)
collection = client.get_collection(name=collection_name)
document_count = collection.count()
if document_count:
    if os.path.exists(FAILED_FILE):
        ids = delete_corrupted_books()
        index = book_ids.index(ids[0])
        # Store mulai dari yang rusak
        book_ids = book_ids[index:]
    else:
        book_ids = []
if book_ids:
    print(
        f"File-file yang ada di direktori {category_dir} akan disimpan dalam "
        f"format vektor di direktori {chroma_dir} ...")
    data = create_data()
    show_sample(data["docs"])
    p = ChromaProgress(data, option.batch_size)
    p.run()

# Inisialisasi ranker (otomatis mengunduh model jika belum ada)
print("Siapkan ranker ...")
ranker = Ranker(model_name=RANKER_MODEL, cache_dir=RANKER_DIR)

print(f"Siapkan pengambil naskah ...")
base_retriever = vector_store.as_retriever(
    search_kwargs=dict(k=option.related_doc_count))

index = option.category_id - 1
category_name = CATEGORIES[index]
print(f"Kategori: {category_name}")
book_count = get_book_count()
print(f"Jumlah buku: {book_count}")

if option.prompt:
    prompt = option.prompt
else:
    prompt = None
    print("Silakan bertanya.")

history = []
while True:
    if not prompt:
        prompt = input(">>> ")
    if prompt == "exit":
        break
    messages = list(history)
    raw_docs = base_retriever.invoke(prompt)
    if raw_docs:
        print(f"Ada {len(raw_docs)} naskah terkait.")
    else:
        print("Tidak ditemukan naskah terkait.")
        break
    # Urutkan yang terbaik
    passages = [
        {"id": idx, "text": d.page_content, "meta": d.metadata}
        for idx, d in enumerate(raw_docs)]
    rerank_request = RerankRequest(query=prompt, passages=passages)
    try:
        rerank_results = ranker.rerank(rerank_request)
    except ValueError:
        # Pasang patch ke ranker session
        ranker.session.run = patched_run
        rerank_results = ranker.rerank(rerank_request)
    # Ambil yang terbaik
    top_results = rerank_results[:option.best_doc_count]
    if top_results:
        print(f"Gunakan {len(top_results)} naskah paling terkait.")
    # Gabungkan teks dokumen beserta ID Buku dan Judul Bab ke dalam konteks LLM
    context_list = []
    for r in top_results:
        text = r["text"]
        book_id = r["meta"]["book_id"]
        page_num = r["meta"]["page_num"]
        citation = f"BUKU {book_id}"
        if page_num:
            citation = f"{citation}, HALAMAN {page_num}"
        text = f"[{citation}]\n{text}"
        context_list.append(text)
    context = "\n\n====================\n\n".join(context_list)
    context = SYSTEM_PROMPT_TEMPLATE.format(context=context)
    messages.append(SystemMessage(context))
    messages.append(SYSTEM_PROMPT_DONT_KNOW)
    messages.append(HumanMessage(prompt))
    history.append(HumanMessage(prompt))
    print(f"Sampaikan prompt ke LLM ...")
    begin_time = time()
    if option.prompt:
        response = llm_client.invoke(messages)
        answer = response.content
        print(f"\n\n{answer}")
        meta = response.response_metadata
    else:
        answer = ""
        for chunk in llm_client.stream(messages):
            print(chunk.content, end="", flush=True)
            answer += chunk.content
            if chunk.usage_metadata:
                meta = chunk.usage_metadata
        print()
    duration = time() - begin_time
    duration = int(duration)
    token = get_token(meta)
    print(
        f"\n\nInput {token['input']} token, Output {token['output']} token, "
        f"{duration} detik")
    history.append(AIMessage(answer))
    if option.prompt:
        break
    prompt = None