Update app.py

app.py

@@ -1,52 +1,104 @@
+# app.py
 import torch
 from transformers import AutoTokenizer, EncoderDecoderModel
+import gradio as gr
 
+# デバイス設定 (Spacesのハードウェア設定に依存)
+# SpacesでGPUを利用する場合、自動的にCUDAが利用可能になります
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}") # デバイス確認用ログ
+
 model_name = "Shuu12121/CodeEncoderDecodeerModel-Ghost"
+print(f"Loading model: {model_name}") # モデル読み込み開始ログ
 
-# Tokenizerの読み込み
-encoder_tokenizer = AutoTokenizer.from_pretrained(f"{model_name}/encoder_tokenizer")
-decoder_tokenizer = AutoTokenizer.from_pretrained(f"{model_name}/decoder_tokenizer")
+# --- Tokenizerの読み込み ---
+try:
+    encoder_tokenizer = AutoTokenizer.from_pretrained(f"{model_name}/encoder_tokenizer")
+    decoder_tokenizer = AutoTokenizer.from_pretrained(f"{model_name}/decoder_tokenizer")
+    print("Tokenizers loaded successfully.")
+except Exception as e:
+    print(f"Error loading tokenizers: {e}")
+    # エラーが発生した場合、Gradioインターフェースでエラーを表示するなどの処理を追加できます
+    raise # ここではエラーを再発生させて、起動を停止させます
 
+# decoder_tokenizerのpad_token設定
 if decoder_tokenizer.pad_token is None:
-    decoder_tokenizer.pad_token = decoder_tokenizer.eos_token
+    if decoder_tokenizer.eos_token is not None:
+        decoder_tokenizer.pad_token = decoder_tokenizer.eos_token
+        print("Set decoder pad_token to eos_token.")
+    else:
+        # eos_tokenもない場合の代替処理（例: '<pad>'トークンを追加）
+        decoder_tokenizer.add_special_tokens({'pad_token': '<pad>'})
+        print("Added '<pad>' as pad_token.")
+        # モデルのリサイズが必要になる場合がある
+        # model.resize_token_embeddings(len(decoder_tokenizer)) # 必要に応じて
 
-model = EncoderDecoderModel.from_pretrained(model_name).to(device)
-model.eval()
+# --- モデルの読み込み ---
+try:
+    model = EncoderDecoderModel.from_pretrained(model_name).to(device)
+    model.eval() # 評価モードに設定
+    print("Model loaded successfully and moved to device.")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    raise
 
+# --- Docstring生成関数 ---
 def generate_docstring(code: str) -> str:
-    inputs = encoder_tokenizer(
-        code,
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-        max_length=2048
-    ).to(device)
-
-    with torch.no_grad():
-        output_ids = model.generate(
-            input_ids=inputs.input_ids,
-            attention_mask=inputs.attention_mask,
-            max_length=256,
-            num_beams=5,
-            early_stopping=True,
-            decoder_start_token_id=model.config.decoder_start_token_id,
-            eos_token_id=model.config.eos_token_id,
-            pad_token_id=model.config.pad_token_id,
-            no_repeat_ngram_size=2
-        )
-
-    return decoder_tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
-# Gradio UI
-import gradio as gr
+    print("Received code snippet for docstring generation.") # 関数呼び出しログ
+    if not code:
+        return "Please provide a code snippet."
+
+    try:
+        # エンコーダー入力の準備
+        inputs = encoder_tokenizer(
+            code,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=2048 # モデルが許容する最大長に合わせる（必要なら調整）
+        ).to(device)
+
+        print(f"Input tokens length: {inputs.input_ids.shape[1]}")
+
+        # 生成実行
+        with torch.no_grad():
+            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,
+                eos_token_id=decoder_tokenizer.eos_token_id, # EOSトークンID
+                pad_token_id=decoder_tokenizer.pad_token_id, # PADトークンID
+                no_repeat_ngram_size=2    # 繰り返さないN-gramサイズ
+            )
+
+        print(f"Generated output tokens length: {output_ids.shape[1]}")
+
+        # デコードしてテキストに変換
+        generated_docstring = decoder_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        print("Docstring generated successfully.")
+        return generated_docstring
+
+    except Exception as e:
+        print(f"Error during generation: {e}")
+        # ユーザーにエラーを通知
+        return f"An error occurred during generation: {e}"
 
+# --- Gradio UI ---
 iface = gr.Interface(
     fn=generate_docstring,
-    inputs=gr.Textbox(label="Code Snippet", lines=10, placeholder="Paste your function here..."),
+    inputs=gr.Textbox(label="Code Snippet", lines=10, placeholder="Paste your Python function or code block here..."),
     outputs=gr.Textbox(label="Generated Docstring"),
-    title="Code-to-Docstring Generator",
-    description="This demo uses a custom encoder-decoder model to generate docstrings from code."
+    title="Code-to-Docstring Generator (Shuu12121/CodeEncoderDecodeerModel-Ghost)",
+    description="This demo uses the Shuu12121/CodeEncoderDecodeerModel-Ghost model to automatically generate Python docstrings from code snippets. Paste your code below and click 'Submit'."
 )
 
-iface.launch()
+# --- アプリケーションの起動 ---
+# Hugging Face Spacesで実行する場合、share=Trueは不要
+if __name__ == "__main__":
+    print("Launching Gradio interface...")
+    iface.launch()
+    print("Gradio interface launched.")