Views
No views yet
1from tqdm import tqdm
2import json
3import os
4
5MODEL_DIR = os.path.join(BASE_DIR, "fine_tuned_model")
6
7
8def generate_predictions(model, tokenizer, input_file, output_file):
9 # バッチ処理の追加
10 BATCH_SIZE = 16 # バッチサイズの設定
11
12 print(f"入力ファイルを読み込み中: {input_file}")
13 tasks = []
14 with open(input_file, 'r', encoding='utf-8') as f:
15 for line in f:
16 tasks.append(json.loads(line))
17
18 results = []
19 print("推論を実行中...")
20
21 # バッチ処理
22 for i in tqdm(range(0, len(tasks), BATCH_SIZE)):
23 batch_tasks = tasks[i:i + BATCH_SIZE]
24 prompts = [f"入力: {task['input']}\n出力: " for task in batch_tasks]
25
26 # バッチでの推論
27 inputs = tokenizer(
28 prompts,
29 return_tensors="pt",
30 padding=True,
31 truncation=True,
32 max_length=512
33 )
34
35 with torch.no_grad():
36 outputs = model.generate(
37 inputs.input_ids,
38 max_length=512,
39 temperature=0.9,
40 do_sample=False,
41 repetition_penalty=1.2,
42 pad_token_id=tokenizer.pad_token_id,
43 top_k=50,
44 top_p=0.95,
45 early_stopping=True, # 早期停止を有効化
46 use_cache=True # キャッシュを使用
47 )
48
49 # バッチ出力の処理
50 for k, task in enumerate(batch_tasks): # 各タスクについてループ
51 output_index = k # インデックスはタスクごとに1つだけ
52 if output_index < len(outputs): # 範囲外アクセスを防ぐ
53 generated_text = tokenizer.decode(outputs[output_index], skip_special_tokens=True)
54 output_text = generated_text.split("出力: ")[-1].strip()
55 results.append({
56 "task_id": task["task_id"], # 正しいタスクIDを取得
57 "output": output_text # 対応する出力
58 })
59
60 print(f"結果を保存中: {output_file}")
61 with open(output_file, 'w', encoding='utf-8') as f:
62 for result in results:
63 json.dump(result, f, ensure_ascii=False)
64 f.write('\n')
65
66def main():
67 # GPUメモリのクリア
68 torch.cuda.empty_cache()
69
70 # 時間計測の追加
71 import time
72 start_time = time.time()
73
74 model, tokenizer = load_model()
75 input_file = "{$file_path}"
76 output_file = os.path.join(BASE_DIR, "{$file_path}")
77
78 generate_predictions(model, tokenizer, input_file, output_file)
79
80 # 実行時間の表示
81 elapsed_time = time.time() - start_time
82 print(f"総実行時間: {elapsed_time / 60:.2f}分")
83
84if __name__ == "__main__":
85 main()