|
|
import os |
|
|
from dataclasses import dataclass |
|
|
from typing import Dict, List |
|
|
|
|
|
@dataclass |
|
|
class Example: |
|
|
user_prompt: str |
|
|
answers_by_model: Dict[str, str] |
|
|
|
|
|
def load_model_responses(base_dir: str) -> List[Example]: |
|
|
|
|
|
model_dirs = [d for d in os.listdir(base_dir) if os.path.isdir(os.path.join(base_dir, d))] |
|
|
|
|
|
prompt_data = {} |
|
|
for model_name in model_dirs: |
|
|
model_path = os.path.join(base_dir, model_name) |
|
|
for fname in sorted(os.listdir(model_path)): |
|
|
if not fname.endswith(".txt"): |
|
|
continue |
|
|
|
|
|
prompt_idx = int(fname.replace("Prompt", "").replace(".txt", "")) |
|
|
|
|
|
fpath = os.path.join(model_path, fname) |
|
|
with open(fpath, "r", encoding="utf-8") as f: |
|
|
lines = [line.strip() for line in f if line.strip()] |
|
|
|
|
|
with open(fpath, "r", encoding="utf-8") as f: |
|
|
text = f.read() |
|
|
|
|
|
turns = [] |
|
|
current_speaker, buffer = None, [] |
|
|
|
|
|
for line in text.splitlines(): |
|
|
if line.startswith("User: "): |
|
|
if current_speaker and buffer: |
|
|
turns.append((current_speaker, "\n".join(buffer).strip())) |
|
|
current_speaker = "User" |
|
|
buffer = [line.replace("User: ", "", 1).strip()] |
|
|
elif line.startswith("Bot: "): |
|
|
if current_speaker and buffer: |
|
|
turns.append((current_speaker, "\n".join(buffer).strip())) |
|
|
current_speaker = "Bot" |
|
|
buffer = [line.replace("Bot: ", "", 1).strip()] |
|
|
else: |
|
|
if current_speaker: |
|
|
buffer.append(line.strip()) |
|
|
|
|
|
if current_speaker and buffer: |
|
|
turns.append((current_speaker, "\n".join(buffer).strip())) |
|
|
|
|
|
if len(turns) < 4: |
|
|
print(f"⚠️ Skipping {fpath} (not enough exchanges)") |
|
|
continue |
|
|
|
|
|
user_prompt = turns[2][1] |
|
|
bot_answer = turns[3][1] |
|
|
|
|
|
if prompt_idx not in prompt_data: |
|
|
prompt_data[prompt_idx] = {"user_prompt": user_prompt, "answers_by_model": {}} |
|
|
|
|
|
prompt_data[prompt_idx]["answers_by_model"][model_name] = bot_answer |
|
|
|
|
|
examples = [ |
|
|
Example(user_prompt=data["user_prompt"], answers_by_model=data["answers_by_model"]) |
|
|
for _, data in sorted(prompt_data.items()) |
|
|
] |
|
|
return examples |
|
|
|
|
|
def save_examples_to_txt(examples: List[Example], out_file: str): |
|
|
with open(out_file, "w", encoding="utf-8") as f: |
|
|
for i, ex in enumerate(examples, start=1): |
|
|
f.write(f"Prompt {i}: {ex.user_prompt}\n") |
|
|
for model, ans in ex.answers_by_model.items(): |
|
|
f.write(f" {model}: {ans}\n") |
|
|
f.write("\n" + "-"*60 + "\n\n") |
|
|
print(f"✅ Saved final answers to {out_file}") |
|
|
|
|
|
if __name__ == "__main__": |
|
|
base_dir = "Results/with_rag/" |
|
|
examples = load_model_responses(base_dir) |
|
|
|
|
|
save_examples_to_txt(examples, "Results/final_answers.txt") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
