Views
No views yet
1!pip install -U langchain-community langchain-huggingface vllm triton wandb weave langchain-huggingface langchain-chroma datasets --q
2
3# GitHubをclone
4!git clone https://github.com/y-hiroki-radiotech/llm-final-task.git
5%cd llm-final-task
6
7import os
8import random
9import numpy as np
10import torch
11import pandas as pd
12from vllm import LLM
13from tqdm import tqdm
14import json
15from datasets import
16import time
17
18from custom_few_shot_prompt_template import CustomFewShotPromptTemplate
19
20
21# JSONLファイルを読み込む
22file_path = 'elyza-tasks-100-TV_0.jsonl'
23data = pd.read_json(file_path, lines=True)
24
25# example selector用のデータ
26df = load_dataset("elyza/ELYZA-tasks-100", split="test")
27df = df.to_pandas()
28examples = []
29for row in df.itertuples():
30 examples.append({"input": row.input, "output": row.output})
31
32
33few_shot = CustomFewShotPromptTemplate(examples)
34# few-shot-selector
35few_shot_list = []
36for row in tqdm(data.itertuples(), desc="生成中"):
37 few_shot_list.append(few_shot.format(row.input))
38
39# 一度キャッシュを削除する
40if torch.cuda.is_available():
41 print("Clearing CUDA cache...")
42 torch.cuda.empty_cache()
43 print("CUDA cache cleared.")
44else:
45 print("CUDA is not available on this system.")
46
47# vllmを使う
48model_name = "hiroki-rad/llm-jp-llm-jp-3-13b-128-ft-3000"
49llm = LLM(model=model_name)
50
51# 2回考えるように推論するクラスをインスタンス化
52from two_stage_think import TwoStageThinking
53
54
55thinking_generator = TwoStageThinking(llm)
56
57# 最終的に1回推論の回答を使うことにした
58results = []
59for row, few_shot in tqdm(zip(data.itertuples(), few_shot_list), desc="生成中"):
60 time.sleep(3)
61 first = thinking_generator.generate_complete_response(row, few_shot)
62 results.append(first)
63
64# データの格納
65jsonl_data = []
66for i in range(len(data)):
67 task_id = data.iloc[i]["task_id"] # Access task_id using the index
68 output = results[i]
69
70 # Create a dictionary for each row
71 jsonl_object = {
72 "task_id": task_id,
73 "output": output
74 }
75 jsonl_data.append(jsonl_object)
76
77with open("llm-jp-vllm-second-thinking-output.jsonl", "w", encoding="utf-8") as outfile:
78 for entry in jsonl_data:
79 # Convert task_id to a regular Python integer before dumping
80 entry["task_id"] = int(entry["task_id"])
81 json.dump(entry, outfile, ensure_ascii=False)
82 outfile.write('\n')