Update app.py

app.py

@@ -13,12 +13,12 @@ print(f"Loading model: {model_name}") # モデル読み込み開始ログ
 
 # --- Tokenizerの読み込み ---
 try:
-    encoder_tokenizer = AutoTokenizer.from_pretrained(f"{model_name}/encoder_tokenizer")
-    decoder_tokenizer = AutoTokenizer.from_pretrained(f"{model_name}/decoder_tokenizer")
+    # subfolder引数を使用してサブディレクトリを指定
+    encoder_tokenizer = AutoTokenizer.from_pretrained(model_name, subfolder="encoder_tokenizer")
+    decoder_tokenizer = AutoTokenizer.from_pretrained(model_name, subfolder="decoder_tokenizer")
     print("Tokenizers loaded successfully.")
 except Exception as e:
     print(f"Error loading tokenizers: {e}")
-    # エラーが発生した場合、Gradioインターフェースでエラーを表示するなどの処理を追加できます
     raise # ここではエラーを再発生させて、起動を停止させます
 
 # decoder_tokenizerのpad_token設定
@@ -35,6 +35,8 @@ if decoder_tokenizer.pad_token is None:
 
 # --- モデルの読み込み ---
 try:
+    # モデルの読み込みは通常通りリポジトリ名を指定すればOK
+    # config.jsonが適切に設定されていれば、エンコーダー/デコーダー部分は自動的に読み込まれる
     model = EncoderDecoderModel.from_pretrained(model_name).to(device)
     model.eval() # 評価モードに設定
     print("Model loaded successfully and moved to device.")
@@ -62,16 +64,18 @@ def generate_docstring(code: str) -> str:
 
         # 生成実行
         with torch.no_grad():
+            # pad_token_idを明示的に指定 (重要: Noneでないことを確認)
+            pad_token_id = decoder_tokenizer.pad_token_id if decoder_tokenizer.pad_token_id is not None else decoder_tokenizer.eos_token_id
+
             output_ids = model.generate(
                 input_ids=inputs.input_ids,
                 attention_mask=inputs.attention_mask,
                 max_length=256,           # 生成するDocstringの最大長
                 num_beams=5,              # ビームサーチのビーム数
                 early_stopping=True,      # 早く停止させるか
-                # decoder_start_token_idは通常model.configから自動設定されるが、明示的に指定も可能
-                # decoder_start_token_id=model.config.decoder_start_token_id,
+                # decoder_start_token_idは通常model.configから自動設定される
                 eos_token_id=decoder_tokenizer.eos_token_id, # EOSトークンID
-                pad_token_id=decoder_tokenizer.pad_token_id, # PADトークンID
+                pad_token_id=pad_token_id, # PADトークンID (Noneでないことを保証)
                 no_repeat_ngram_size=2    # 繰り返さないN-gramサイズ
             )