app.py

"""校異源氏物語の類似テキスト検索システム

このモジュールは、校異源氏物語のテキストデータベースに対して
類似テキスト検索を行うWebインターフェースを提供します。
"""

import json
import xml.etree.ElementTree as ET

import gradio as gr
from Levenshtein import ratio

DATA_PATH = "./data.json"

with open(DATA_PATH, "r", encoding="utf-8") as f:
    documents_data = json.load(f)

def predict(query, selected_vols, top_n=5):
    """テキストの類似度を計算し、上位の結果を返す

    Args:
        query (str): 検索クエリテキスト
        selected_vols (list): 検索対象の巻のリスト
        top_n (int, optional): 返す結果の数. デフォルトは5

    Returns:
        list: スコア順にソートされた上位n件の検索結果
    """
    results = []
    
    for doc in documents_data:
        # 選択された巻のみを検索対象とする
        if not selected_vols or str(doc["vol"]) in selected_vols:
            score = ratio(query, doc["text"])
            results.append({
                "vol": doc["vol"],
                "page": doc["page"],
                "score": score,
                "text": doc["text"]
            })

    results.sort(key=lambda x: x["score"], reverse=True)
    top_results = results[:top_n]  # top_nで指定された件数だけを取得

    return top_results


def extract_text_from_lines(element):
    """本文タイプの要素からテキストを抽出する"""
    lines = element.findall(".//*[@type='本文']")
    return ''.join(line.text for line in lines)

def format_prediction_result(result):
    """予測結果を 'vol-page' 形式にフォーマットする"""
    first_result = result[0]
    return f'{first_result["vol"]}-{first_result["page"]}'

def search_similar_texts(query, selected_vols, top_n=5, xml_file=None):
    """テキストの類似検索を実行する関数

    Args:
        query (str): 検索クエリテキスト
        selected_vols (list): 検索対象の巻のリスト
        top_n (int, optional): 返す結果の数. デフォルトは5
        xml_file (gradio.File, optional): 比較対象のXMLファイル

    Returns:
        list: 検索結果のリスト。XMLファイル処理時は[predict_results]、
             通常検索時は[top_results]を返す
    """
    if xml_file is not None:
        try:
            # Gradioのファイルオブジェクトから名前を取得して直接ファイルを開く
            xml_content = xml_file.name
            tree = ET.parse(xml_content)
            root = tree.getroot()
            
            # ページ要素の取得
            elements = root.findall(".//*[@type='page']")

            # 予測実行
            predict_results = {}
            for i, element in enumerate(elements, 1):  # enumerate(elements, 1)で1から開始
                text = extract_text_from_lines(element)
                top_results = predict(text, selected_vols, 1)
                predict_results[str(i)] = format_prediction_result(top_results)

            return [predict_results]
    
        except (ET.ParseError, FileNotFoundError, PermissionError) as e:
            print(f"XMLファイルの処理中にエラーが発生しました: {str(e)}")
            return [[], {}]
    

    top_results = predict(query, selected_vols, top_n)
    
    return [top_results] # , vol_percentages

# Gradioインターフェースの作成
demo = gr.Interface(
    fn=search_similar_texts,
    inputs=[
        gr.Textbox(label="検索テキスト", placeholder="検索したいテキストを入力してください"),
        gr.Dropdown(
            choices=[str(i) for i in range(1, 55)],
            label="巻",
            multiselect=True,
            value=[],
        ),
        gr.Slider(minimum=1, maximum=10, value=5, step=1, label="表示件数"),
        gr.File(label="TEI/XMLファイル")
    ],
    outputs=[
        gr.JSON(),
        # gr.JSON(label="巻ごとの割合")
    ],
    title="校異源氏物語 類似テキスト検索",
    description="テキストを入力すると、校異源氏物語の類似する箇所を検索します。TEI/XMLファイルのアップロードにあたっては、[こちら](https://zenn.dev/nakamura196/articles/e9238881bbc4f7#ocr)をご覧ください。",
    allow_flagging="never",
    examples=[
        ["同五源氏誕生より十二才まて有いづれの御時にか。女御更衣ありさふらひ浴けるなかにいとやむこ□□□となききはにはあらぬかすぐれてときめき給ふありけりはしめよりわれはと思ひあがり給へる御かためざましきものにおとしめそねみ給おなじほどそれより下らうの更衣たちはましてやすからすあさ夕の宮つかへにつけても人の心をうこかしうらみをおふつもりにや有けん銅大世更衣いとあづしくなりゆきもの心ほそけにさとがちなるをいよあかすあつれなるものにおほゝして人のそしりをもえはゞからせ給はす世のためしにもなりぬべき御もてなしなりかんだちめうへ人などもあひなくめを", [], 5, None]
    ]
)

# インターフェースの起動
demo.launch()