final version

README.md

@@ -7,7 +7,7 @@ base_model:
 - llm-jp/llm-jp-3-13b
 ---
 ```python
-!pip install -U transformers vllm triton --q
+!pip install -U langchain-community langchain-huggingface vllm triton wandb weave langchain-huggingface langchain-chroma datasets --q
 
 # GitHubをclone
 !git clone https://github.com/y-hiroki-radiotech/llm-final-task.git
@@ -21,23 +21,36 @@ import pandas as pd
 from vllm import LLM
 from tqdm import tqdm
 import json
+from datasets import 
+
+from custom_few_shot_prompt_template import CustomFewShotPromptTemplate
+
 
 # JSONLファイルを読み込む
 file_path = 'elyza-tasks-100-TV_0.jsonl'
 data = pd.read_json(file_path, lines=True)
 
+# example selector用のデータ
+df = load_dataset("elyza/ELYZA-tasks-100", split="test")
+df = df.to_pandas()
+examples = []
+for row in df.itertuples():
+    examples.append({"input": row.input, "output": row.output})
 
-def set_seed(seed):
-   random.seed(seed)
-   os.environ["PYTHONHASHSEED"] = str(seed)
-   np.random.seed(seed)
-   torch.manual_seed(seed)
-   torch.cuda.manual_seed(seed)
-   torch.cuda.manual_seed_all(seed)
-   torch.backends.cudnn.deterministic = True
-   torch.backends.cudnn.benchmark = False
 
-set_seed(42)
+few_shot = CustomFewShotPromptTemplate(examples)
+# few-shot-selector
+few_shot_list = []
+for row in tqdm(data.itertuples(), desc="生成中"):
+    few_shot_list.append(few_shot.format(row.input))
+
+# 一度キャッシュを削除する
+if torch.cuda.is_available():
+    print("Clearing CUDA cache...")
+    torch.cuda.empty_cache()
+    print("CUDA cache cleared.")
+else:
+    print("CUDA is not available on this system.")
 
 # vllmを使う
 model_name = "hiroki-rad/llm-jp-llm-jp-3-13b-128-ft-3000"
@@ -46,18 +59,17 @@ llm = LLM(model=model_name)
 # 2回考えるように推論するクラスをインスタンス化
 from two_stage_think import TwoStageThinking
 
+
 thinking_generator = TwoStageThinking(llm)
 
-# データフレームでの使用例
+# 最終的に1回推論の回答を使うことにした
 results = []
-for row in tqdm(data.itertuples(), desc="生成中"):
-    first, second = thinking_generator.generate_complete_response(row)
-    results.append(second)
-
+for row, few_shot in tqdm(zip(data.itertuples(), few_shot_list), desc="生成中"):
+    first = thinking_generator.generate_complete_response(row, few_shot)
+    results.append(first)
 
+# データの格納
 jsonl_data = []
-
-# Iterate through the data and outputs
 for i in range(len(data)):
     task_id = data.iloc[i]["task_id"] # Access task_id using the index
     output = results[i]