Update app.py

app.py

@@ -4,19 +4,28 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import spaces
 
-# PyTorch settings for better performance and determinism
+# PyTorch設定（パフォーマンスと再現性向上のため）
 torch.backends.cudnn.deterministic = True
 torch.backends.cudnn.benchmark = False
 torch.backends.cuda.matmul.allow_tf32 = True
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# トークナイザーとモデルのロード
-tokenizer = AutoTokenizer.from_pretrained('Local-Novel-LLM-project/Vecteus-v1-abliterated', attn_implementation="flash_attention_2", use_auth_token=HF_TOKEN)
-model = AutoModelForCausalLM.from_pretrained('Local-Novel-LLM-project/Vecteus-v1-abliterated', use_auth_token=HF_TOKEN)
+# モデルのキャッシュ用辞書（ロード済みなら再利用）
+loaded_models = {}
+
+def get_model_and_tokenizer(model_name):
+    if model_name in loaded_models:
+        return loaded_models[model_name]
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, attn_implementation="flash_attention_2", use_auth_token=HF_TOKEN)
+        model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=HF_TOKEN)
+        loaded_models[model_name] = (model, tokenizer)
+        return model, tokenizer
 
 @spaces.GPU
 def generate_text(
+    model_name,
     input_text, 
     max_length=150, 
     temperature=0.7, 
@@ -29,6 +38,9 @@ def generate_text(
         if not input_text.strip():
             return ""
 
+        # 選択されたモデルとトークナイザーを取得
+        model, tokenizer = get_model_and_tokenizer(model_name)
+        
         # GPUが利用可能ならGPUへ移動。bf16がサポートされている場合はbf16を使用
         device = "cuda" if torch.cuda.is_available() else "cpu"
         if device == "cuda" and hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported():
@@ -36,11 +48,11 @@ def generate_text(
         else:
             model.to(device)
 
-        # 入力テキストをトークン化してトークン数を取得
+        # 入力テキストのトークン化
         input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
         input_token_count = input_ids.shape[1]
         
-        # 総トークン数の上限を入力トークン数 + max_length（max_lengthはトークン数として扱う）
+        # 総トークン数の上限を入力トークン数 + max_length（max_lengthは追加するトークン数として扱う）
         total_max_length = input_token_count + max_length
         
         # テキスト生成
@@ -56,7 +68,7 @@ def generate_text(
             num_return_sequences=1
         )
         
-        # 生成されたテキストをデコードし、入力部分を除いた部分を抽出
+        # 生成されたテキストをデコードし、入力部分を除いた生成分を抽出
         generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
         new_text = generated_text[len(input_text):]
         
@@ -65,29 +77,40 @@ def generate_text(
     except Exception as e:
         return f"{input_text}\n\nエラーが発生しました: {str(e)}"
 
-# Gradio インターフェースの作成
+# Gradioインターフェースの作成
 with gr.Blocks() as demo:
     gr.Markdown("# テキスト続き生成アシスタント")
-    gr.Markdown("テキストボックスに文章を入力し、パラメータを調整して「続きを生成」ボタンをクリックすると、AIがその続きを現在のテキストに追加します。")
+    gr.Markdown("モデルを選択し、テキストボックスに文章を入力してパラメータを調整後、「続きを生成」ボタンをクリックすると、選択したモデルがその続きを生成します。")
+    
+    # モデル選択用プルダウンメニュー（候補は必要に応じて変更してください）
+    model_dropdown = gr.Dropdown(
+        choices=[
+            "Local-Novel-LLM-project/Vecteus-v1-abliterated", 
+            "gpt2", 
+            "EleutherAI/gpt-neo-125M"
+        ],
+        label="モデルを選択してください",
+        value="Local-Novel-LLM-project/Vecteus-v1-abliterated"
+    )
     
-    # 単一のテキストボックス
+    # テキスト入力ボックス
     input_text = gr.Textbox(label="テキストを入力してください", placeholder="ここにテキストを入力...", lines=10)
     
-    # パラメータの設定UI - すべて縦に一列で配置
+    # 生成パラメータの設定UI（縦に一列で配置）
     max_length_slider = gr.Slider(minimum=10, maximum=1000, value=100, step=10, label="追加するトークン数")
     temperature_slider = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="創造性（温度）")
     top_k_slider = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="top_k")
     top_p_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.05, label="top_p")
     repetition_penalty_slider = gr.Slider(minimum=1.0, maximum=2.0, value=1.0, step=0.1, label="繰り返しペナルティ")
     
-    # ボタンも縦に配置
+    # ボタン（縦に配置）
     generate_btn = gr.Button("続きを生成", variant="primary")
     clear_btn = gr.Button("クリア")
     
-    # イベントの設定
+    # イベントの設定：入力としてモデル選択とテキスト、パラメータを渡す
     generate_btn.click(
         fn=generate_text, 
-        inputs=[input_text, max_length_slider, temperature_slider, top_k_slider, top_p_slider, repetition_penalty_slider], 
+        inputs=[model_dropdown, input_text, max_length_slider, temperature_slider, top_k_slider, top_p_slider, repetition_penalty_slider], 
         outputs=input_text
     )
     
@@ -96,17 +119,18 @@ with gr.Blocks() as demo:
     # 使い方の説明
     gr.Markdown("""
     ## 使い方
-    1. テキストボックスに続きを生成したい文章を入力します
-    2. 生成パラメータ（追加するトークン数、創造性、top_k、top_p、繰り返しペナルティ）を調整します
-    3. 「続きを生成」ボタンをクリックすると、入力したテキストの続きが生成され、元のテキストに追加されます
-    4. 「クリア」ボタンを押すと、テキストボックスの内容がクリアされます
-    5. 満足のいく結果が得られるまで、繰り返し「続きを生成」ボタンを押して文章を発展させることができます
+    1. 上部のプルダウンメニューから使用するモデルを選択します
+    2. テキストボックスに続きを生成したい文章を入力します
+    3. 生成パラメータ（追加するトークン数、創造性、top_k、top_p、繰り返しペナルティ）を調整します
+    4. 「続きを生成」ボタンをクリックすると、入力したテキストの続きが生成され、元のテキストに追加されます
+    5. 「クリア」ボタンを押すと、テキストボックスの内容がクリアされます
+    6. 満足のいく結果が得られるまで、繰り返し「続きを生成」ボタンを押して文章を発展させることができます
     
     ## ヒント
-    - 短い文章から始めて徐々に発展させると良い結果が得られます
+    - 短い文章から始め、徐々に発展させると良い結果が得られます
     - 創造性（温度）を高くすると予測不可能な生成結果に、低くすると安定した結果になります
-    - top_k や top_p、繰り返しペナルティなども状況に応じて調整してみてください
+    - top_k や top_p、繰り返しペナルティも状況に応じて調整してみてください
     """)
 
 # アプリの起動
-demo.launch()
+demo.launch()