update Readme

README.md

@@ -14,24 +14,16 @@ tags:
 
 # Sample Use
 ```python
-MODEL_DIR = os.path.join("model_dir")
+from tqdm import tqdm
+import json
+import os
 
-def load_model():
-    print("モデルとトークナイザーを読み込み中...")
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_DIR,
-        torch_dtype=torch.float16,
-        device_map={"": 0},  # 明示的にGPU割り当て
-        use_cache=True,      # キャッシュを有効化
-    ).to('cuda')            # 明示的にGPUへ
+MODEL_DIR = os.path.join(BASE_DIR, "fine_tuned_model")
 
-    model.eval()  # 評価モード
-    return model, tokenizer
 
 def generate_predictions(model, tokenizer, input_file, output_file):
     # バッチ処理の追加
-    BATCH_SIZE = 8  # バッチサイズの設定
+    BATCH_SIZE = 16  # バッチサイズの設定
 
     print(f"入力ファイルを読み込み中: {input_file}")
     tasks = []
@@ -54,30 +46,32 @@ def generate_predictions(model, tokenizer, input_file, output_file):
             padding=True,
             truncation=True,
             max_length=512
-        ).to('cuda')
+        )
 
         with torch.no_grad():
             outputs = model.generate(
                 inputs.input_ids,
                 max_length=512,
-                temperature=0.7,
+                temperature=0.9,
                 do_sample=False,
                 repetition_penalty=1.2,
                 pad_token_id=tokenizer.pad_token_id,
-                num_return_sequences=1,
+                top_k=50,
+                top_p=0.95,
                 early_stopping=True,    # 早期停止を有効化
                 use_cache=True          # キャッシュを使用
             )
 
         # バッチ出力の処理
-        for j, output in enumerate(outputs):
-            generated_text = tokenizer.decode(output, skip_special_tokens=True)
-            output_text = generated_text.split("出力: ")[-1].strip()
-
-            results.append({
-                "task_id": batch_tasks[j]["task_id"],
-                "output": output_text
-            })
+        for k, task in enumerate(batch_tasks):  # 各タスクについてループ
+          output_index = k  # インデックスはタスクごとに1つだけ
+          if output_index < len(outputs):  # 範囲外アクセスを防ぐ
+              generated_text = tokenizer.decode(outputs[output_index], skip_special_tokens=True)
+              output_text = generated_text.split("出力: ")[-1].strip()
+              results.append({
+                  "task_id": task["task_id"],  # 正しいタスクIDを取得
+                  "output": output_text       # 対応する出力
+              })
 
     print(f"結果を保存中: {output_file}")
     with open(output_file, 'w', encoding='utf-8') as f: