Update app.py

app.py

@@ -3,6 +3,7 @@ import re
 import glob
 import csv
 import threading
+import shutil
 from datetime import datetime, timedelta
 from typing import Optional
 
@@ -14,20 +15,19 @@ from openai import OpenAI
 # 設定（元コード踏襲）
 # =========================
 API_KEY = os.getenv("API_KEY")
-BASE_URL = os.getenv("BASE_URL", "https://openrouter.ai/api/v1")
+BASE_URL = "https://openrouter.ai/api/v1"
 MODEL = os.getenv("MODEL", "google/gemini-2.5-flash")
 
-# Hugging Face Spaces 永続ストレージ（Persistent Storage 有効なら /data が永続）
-OUT_DIR = os.getenv("OUT_DIR", "/data")
+# Free Space では /data は永続ではない可能性が高いので、
+# ダウンロード前提で「カレント」に保存してOK
+OUT_DIR = os.getenv("OUT_DIR", ".")
 os.makedirs(OUT_DIR, exist_ok=True)
 
-# スコアCSV（任意だが便利）
 CSV_PATH = os.path.join(OUT_DIR, "rewrite_scores.csv")
 
 _word_re = re.compile(r"[A-Za-z]+(?:'[A-Za-z]+)?")
 
-# ★追加：書き換え本文をtxtで蓄積するフォルダ（要件：rewrite_passages）
-# 既定は /data/rewrite_passages（永続に残る）
+# 書き換え本文をtxtで蓄積するフォルダ
 REWRITE_DIR = "rewrite_passages"
 os.makedirs(REWRITE_DIR, exist_ok=True)
 
@@ -38,7 +38,6 @@ if not API_KEY:
 
 client = OpenAI(base_url=BASE_URL, api_key=API_KEY)
 
-# 同時実行を軽く制限（Spacesで安定させる）
 REWRITE_CONCURRENCY = int(os.getenv("REWRITE_CONCURRENCY", "2"))
 _rewrite_sem = threading.Semaphore(REWRITE_CONCURRENCY)
 
@@ -83,10 +82,9 @@ excluding the title, author name, source information, chapter number, annotation
 {text}
 """
 
-    # 安全な max_tokens 候補（大→小）
     max_tokens_candidates = [3000, 2000, 1500, 1000]
-
     last_error = None
+
     for mt in max_tokens_candidates:
         try:
             with _rewrite_sem:
@@ -94,22 +92,21 @@ excluding the title, author name, source information, chapter number, annotation
                     model=MODEL,
                     messages=[{"role": "user", "content": prompt}],
                     temperature=0.4,
-                    max_tokens=mt
+                    max_tokens=mt,
                 )
             return resp.choices[0].message.content.strip()
 
         except Exception as e:
             last_error = e
-            # 402（クレジット不足 or トークン過多）の場合は縮めて再試行
             if "402" in str(e) or "more credits" in str(e):
                 continue
-            else:
-                raise e
+            raise
 
-    # 全部ダメなら最後のエラーを投げる
     raise RuntimeError(f"Rewrite failed after retries: {last_error}")
 
-
+# =========================
+# 指標（FRE + 単語数）
+# =========================
 def count_words_english(text: str) -> int:
     return len(_word_re.findall(text))
 
@@ -119,12 +116,11 @@ def compute_metrics(text: str) -> tuple[float, int]:
     return fre, wc
 
 # =========================
-# CSV追記（軽量・永続）
+# CSV追記
 # =========================
 _csv_lock = threading.Lock()
 
 def append_csv_row(row: dict):
-    """rewrite_scores.csv に1行追記（ヘッダ無ければ作成）"""
     fieldnames = ["timestamp_jst", "Text#", "target_level", "flesch_reading_ease", "word_count", "rewritten_text"]
     with _csv_lock:
         exists = os.path.exists(CSV_PATH)
@@ -135,7 +131,7 @@ def append_csv_row(row: dict):
             w.writerow({k: row.get(k, "") for k in fieldnames})
 
 # =========================
-# ★追加：rewrite_passages に txt 追記
+# txt追記
 # =========================
 _txt_lock = threading.Lock()
 
@@ -146,9 +142,6 @@ def append_rewrite_txt(
     word_count: int,
     rewritten_text: str,
 ):
-    """
-    rewrite_passages/Text_{id}.txt に結果を追記（同じText#の再実行も追記で残す）
-    """
     ts = (datetime.utcnow() + timedelta(hours=9)).strftime("%Y-%m-%d %H:%M:%S")
     path = os.path.join(REWRITE_DIR, f"Text_{text_id}.txt")
 
@@ -158,11 +151,9 @@ def append_rewrite_txt(
         f"Flesch Reading Ease: {fre:.2f}\n"
         f"Word Count: {word_count}\n"
         f"Timestamp (JST): {ts}\n"
-        f"Model: {MODEL}\n"
-        f"\n"
+        f"Model: {MODEL}\n\n"
         f"---- Rewritten Text ----\n"
-        f"{rewritten_text}\n"
-        f"\n"
+        f"{rewritten_text}\n\n"
         f"{'=' * 80}\n"
     )
 
@@ -194,10 +185,10 @@ def start(level: int):
     st = init_state()
     total = len(st["files"])
     if total == 0:
-        return st, "passages/pg*.txt が見つかりません", "", "", "", gr.update(visible=False)
+        return st, "passages/pg*.txt が見つかりません", "", "", "", None, None
 
     msg = f"準備完了: {total}件。次に処理するのは #Text {st['files'][0][0]} です。"
-    return st, msg, "", "", "", gr.update(visible=True)
+    return st, msg, "", "", "", None, None
 
 def run_one(level: int, state: dict):
     set_stop(False)
@@ -207,7 +198,7 @@ def run_one(level: int, state: dict):
     total = len(files)
 
     if idx >= total:
-        return state, "全て処理済みです。", "", "", "", gr.update(visible=True)
+        return state, "全て処理済みです。", "", "", "", None, None
 
     text_id, path = files[idx]
     original = load_text(path)
@@ -215,10 +206,8 @@ def run_one(level: int, state: dict):
     rewritten = rewrite_level(original, target_level=level)
     fre, wc = compute_metrics(rewritten)
 
-    # JSTタイムスタンプ（+9）
     ts = (datetime.utcnow() + timedelta(hours=9)).strftime("%Y-%m-%d %H:%M:%S")
 
-    # CSV追記
     append_csv_row({
         "timestamp_jst": ts,
         "Text#": text_id,
@@ -228,7 +217,6 @@ def run_one(level: int, state: dict):
         "rewritten_text": rewritten
     })
 
-    # ★txt追記
     append_rewrite_txt(
         text_id=text_id,
         target_level=level,
@@ -248,10 +236,9 @@ def run_one(level: int, state: dict):
         f"Saved TXT: {os.path.join(REWRITE_DIR, f'Text_{text_id}.txt')}"
     )
     progress = f"{state['idx']} / {total}"
-    return state, "1件処理しました。", header, progress, rewritten, gr.update(visible=True)
+    return state, "1件処理しました。", header, progress, rewritten, None, None
 
 def run_all(level: int, state: dict):
-    """全件（または残り）を順次処理。途中で「停止」ボタンで止められる。"""
     set_stop(False)
 
     files = state.get("files", [])
@@ -259,7 +246,7 @@ def run_all(level: int, state: dict):
     total = len(files)
 
     if idx >= total:
-        return state, "全て処理済みです。", "", f"{idx} / {total}", "", gr.update(visible=True)
+        return state, "全て処理済みです。", "", f"{idx} / {total}", "", None, None
 
     last_header = ""
     last_text = ""
@@ -267,7 +254,7 @@ def run_all(level: int, state: dict):
     while idx < total:
         if get_stop():
             state["idx"] = idx
-            return state, "停止しました。", last_header, f"{idx} / {total}", last_text, gr.update(visible=True)
+            return state, "停止しました。", last_header, f"{idx} / {total}", last_text, None, None
 
         text_id, path = files[idx]
         original = load_text(path)
@@ -286,7 +273,6 @@ def run_all(level: int, state: dict):
             "rewritten_text": rewritten
         })
 
-        # ★txt追記
         append_rewrite_txt(
             text_id=text_id,
             target_level=level,
@@ -308,7 +294,7 @@ def run_all(level: int, state: dict):
         idx += 1
         state["idx"] = idx
 
-    return state, "全件処理が完了しました。", last_header, f"{idx} / {total}", last_text, gr.update(visible=True)
+    return state, "全件処理が完了しました。", last_header, f"{idx} / {total}", last_text, None, None
 
 def stop():
     set_stop(True)
@@ -321,9 +307,6 @@ def reset_csv():
     return f"CSVを削除しました: {CSV_PATH}"
 
 def reset_rewrite_folder():
-    """
-    rewrite_passages を全消しは危険なので、ここでは「中の txt を削除」する実装。
-    """
     removed = 0
     with _txt_lock:
         for fp in glob.glob(os.path.join(REWRITE_DIR, "Text_*.txt")):
@@ -334,6 +317,65 @@ def reset_rewrite_folder():
                 pass
     return f"rewrite_passages の Text_*.txt を削除しました（{removed}件）: {REWRITE_DIR}"
 
+# =========================
+# ★追加：ダウンロード機能
+# =========================
+def list_generated_txt_files() -> list[str]:
+    files = sorted(glob.glob(os.path.join(REWRITE_DIR, "Text_*.txt")))
+    return [os.path.basename(f) for f in files]
+
+def build_single_txt_path(selected_name: str) -> str:
+    """
+    選択されたtxtをダウンロード用に返す（gr.Fileに渡す）
+    """
+    path = os.path.join(REWRITE_DIR, selected_name)
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"Not found: {path}")
+    return path
+
+def build_zip_of_txts(mode: str, n_last: int) -> str:
+    """
+    mode:
+      - 'all' : 全txt
+      - 'last_n' : 最新N個
+    """
+    files = sorted(glob.glob(os.path.join(REWRITE_DIR, "Text_*.txt")), key=os.path.getmtime)
+    if not files:
+        raise FileNotFoundError("No generated txt files yet.")
+
+    if mode == "last_n":
+        files = files[-max(1, int(n_last)):]
+
+    # ZIPは毎回作り直す
+    zip_path = os.path.join(OUT_DIR, "rewrite_passages.zip")
+    if os.path.exists(zip_path):
+        os.remove(zip_path)
+
+    # 一時ディレクトリに集めてからzip（shutil.make_archiveはディレクトリ単位）
+    tmp_dir = os.path.join(OUT_DIR, "_zip_tmp")
+    if os.path.exists(tmp_dir):
+        shutil.rmtree(tmp_dir)
+    os.makedirs(tmp_dir, exist_ok=True)
+
+    for fp in files:
+        shutil.copy(fp, os.path.join(tmp_dir, os.path.basename(fp)))
+
+    shutil.make_archive(os.path.join(OUT_DIR, "rewrite_passages"), "zip", tmp_dir)
+
+    # 後片付け
+    shutil.rmtree(tmp_dir, ignore_errors=True)
+
+    return zip_path
+
+def download_csv() -> str:
+    if not os.path.exists(CSV_PATH):
+        raise FileNotFoundError("rewrite_scores.csv is not created yet.")
+    return CSV_PATH
+
+def refresh_txt_dropdown() -> gr.Dropdown:
+    names = list_generated_txt_files()
+    return gr.Dropdown(choices=names, value=(names[-1] if names else None))
+
 # =========================
 # Gradio UI（Spaces向け）
 # =========================
@@ -358,15 +400,46 @@ with gr.Blocks() as demo:
         reset_btn = gr.Button("CSVリセット（削除）")
         reset_txt_btn = gr.Button("rewrite_passagesリセット（Text_*.txt削除）")
 
-    gr.Markdown(f"📄 CSV保存先: `{CSV_PATH}`（SpacesのFilesに出ます）")
-    gr.Markdown(f"📝 TXT保存先: `{REWRITE_DIR}`（SpacesのFilesに出ます）")
+    gr.Markdown("## 📥 Download")
+    gr.Markdown("Free Space では Files に生成物が見えないことがあるので、ここからダウンロードしてください。")
 
-    start_btn.click(fn=start, inputs=[level], outputs=[state, status, header, progress, output_text])
-    one_btn.click(fn=run_one, inputs=[level, state], outputs=[state, status, header, progress, output_text])
-    all_btn.click(fn=run_all, inputs=[level, state], outputs=[state, status, header, progress, output_text])
+    with gr.Row():
+        refresh_btn = gr.Button("txt一覧を更新")
+        txt_dropdown = gr.Dropdown(choices=list_generated_txt_files(), label="生成済み txt を選択（1個ダウンロード）")
+        download_one_btn = gr.Button("選択したtxtをダウンロード")
+        download_one_file = gr.File(label="Download (single txt)")
+
+    with gr.Row():
+        zip_mode = gr.Radio(
+            choices=[("全txtをまとめてZIP", "all"), ("最新N個をZIP", "last_n")],
+            value="all",
+            label="ZIPの作り方"
+        )
+        last_n = gr.Number(value=5, precision=0, label="N（最新N個の場合）")
+        download_zip_btn = gr.Button("ZIPを作ってダウンロード")
+        download_zip_file = gr.File(label="Download (zip)")
+
+    with gr.Row():
+        download_csv_btn = gr.Button("CSV（rewrite_scores.csv）をダウンロード")
+        download_csv_file = gr.File(label="Download (csv)")
+
+    # 既存の表示（参考）
+    gr.Markdown(f"📄 CSVパス: `{CSV_PATH}`")
+    gr.Markdown(f"📝 TXTフォルダ: `{REWRITE_DIR}`")
+
+    # ---- 既存ボタン ----
+    start_btn.click(fn=start, inputs=[level], outputs=[state, status, header, progress, output_text, download_one_file, download_zip_file])
+    one_btn.click(fn=run_one, inputs=[level, state], outputs=[state, status, header, progress, output_text, download_one_file, download_zip_file])
+    all_btn.click(fn=run_all, inputs=[level, state], outputs=[state, status, header, progress, output_text, download_one_file, download_zip_file])
     stop_btn.click(fn=stop, inputs=[], outputs=[status])
     reset_btn.click(fn=reset_csv, inputs=[], outputs=[status])
     reset_txt_btn.click(fn=reset_rewrite_folder, inputs=[], outputs=[status])
 
+    # ---- Download UI ----
+    refresh_btn.click(fn=refresh_txt_dropdown, inputs=[], outputs=[txt_dropdown])
+    download_one_btn.click(fn=build_single_txt_path, inputs=[txt_dropdown], outputs=[download_one_file])
+    download_zip_btn.click(fn=build_zip_of_txts, inputs=[zip_mode, last_n], outputs=[download_zip_file])
+    download_csv_btn.click(fn=download_csv, inputs=[], outputs=[download_csv_file])
+
 demo.queue(max_size=64)
 demo.launch()