feat: add tei xml option

app.py

@@ -1,13 +1,31 @@
+"""校異源氏物語の類似テキスト検索システム
+
+このモジュールは、校異源氏物語のテキストデータベースに対して
+類似テキスト検索を行うWebインターフェースを提供します。
+"""
+
+import json
+import xml.etree.ElementTree as ET
+
 import gradio as gr
 from Levenshtein import ratio
-import json
 
-data_path = "./data.json"
+DATA_PATH = "./data.json"
 
-with open(data_path, "r") as f:
+with open(DATA_PATH, "r", encoding="utf-8") as f:
     documents_data = json.load(f)
 
-def search_similar_texts(query, selected_vols, top_n=5):
+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:
@@ -24,24 +42,58 @@ def search_similar_texts(query, selected_vols, top_n=5):
     results.sort(key=lambda x: x["score"], reverse=True)
     top_results = results[:top_n]  # top_nで指定された件数だけを取得
 
-    '''
-    # top_n件の結果における巻ごとの件数をカウント
-    vol_counts = {}
-    total_results = len(top_results)
-    
-    for doc in top_results:
-        vol_str = f"巻{doc['vol']}"
-        vol_counts[vol_str] = vol_counts.get(vol_str, 0) + 1
+    return top_results
+
 
-    # 巻ごとの割合を計算
-    vol_percentages = {
-        vol: f"{(count/total_results*100):.1f}%" 
-        for vol, count in vol_counts.items()
-    } if total_results > 0 else {}
+def extract_text_from_lines(element):
+    """本文タイプの要素からテキストを抽出する"""
+    lines = element.findall(".//*[@type='本文']")
+    return ''.join(line.text for line in lines)
 
-    return [top_results, vol_percentages]  # リストとして返す
-    '''
+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インターフェースの作成
@@ -55,19 +107,20 @@ demo = gr.Interface(
             multiselect=True,
             value=[],
         ),
-        gr.Slider(minimum=1, maximum=10, value=5, step=1, label="表示件数")
+        gr.Slider(minimum=1, maximum=10, value=5, step=1, label="表示件数"),
+        gr.File(label="TEI/XMLファイル")
     ],
     outputs=[
         gr.JSON(),
         # gr.JSON(label="巻ごとの割合")
     ],
     title="校異源氏物語 類似テキスト検索",
-    description="テキストを入力すると、校異源氏物語の類似する箇所を検索します。",
+    description="テキストを入力すると、校異源氏物語の類似する箇所を検索します。TEI/XMLファイルのアップロードにあたっては、[こちら](https://zenn.dev/nakamura196/articles/e9238881bbc4f7#ocr)をご覧ください。",
     allow_flagging="never",
     examples=[
-        ["同五源氏誕生より十二才まて有いづれの御時にか。女御更衣ありさふらひ浴けるなかにいとやむこ□□□となききはにはあらぬかすぐれてとき���き給ふありけりはしめよりわれはと思ひあがり給へる御かためざましきものにおとしめそねみ給おなじほどそれより下らうの更衣たちはましてやすからすあさ夕の宮つかへにつけても人の心をうこかしうらみをおふつもりにや有けん銅大世更衣いとあづしくなりゆきもの心ほそけにさとがちなるをいよあかすあつれなるものにおほゝして人のそしりをもえはゞからせ給はす世のためしにもなりぬべき御もてなしなりかんだちめうへ人などもあひなくめを", [], 5]
+        ["同五源氏誕生より十二才まて有いづれの御時にか。女御更衣ありさふらひ浴けるなかにいとやむこ□□□となききはにはあらぬかすぐれてときめき給ふありけりはしめよりわれはと思ひあがり給へる御かためざましきものにおとしめそねみ給おなじほどそれより下らうの更衣たちはましてやすからすあさ夕の宮つかへにつけても人の心をうこかしうらみをおふつもりにや有けん銅大世更衣いとあづしくなりゆきもの心ほそけにさとがちなるをいよあかすあつれなるものにおほゝして人のそしりをもえはゞからせ給はす世のためしにもなりぬべき御もてなしなりかんだちめうへ人などもあひなくめを", [], 5, None]
     ]
 )
 
 # インターフェースの起動
-demo.launch()
+demo.launch()