| | """ |
| | Usage: |
| | python3 summarize_cluster.py --in results_c20_kmeans_cluster.pkl --model gpt-4 --num-prompts 100 |
| | python3 summarize_cluster.py --in results_c20_kmeans_cluster.pkl --model azure-gpt-4-32k --num-prompts 200 |
| | """ |
| | import argparse |
| | import pickle |
| |
|
| | import pandas as pd |
| |
|
| | from fastchat.llm_judge.common import ( |
| | chat_completion_openai, |
| | chat_completion_openai_azure, |
| | chat_completion_anthropic, |
| | ) |
| | from fastchat.conversation import get_conv_template |
| |
|
| |
|
| | def truncate_string(s, l): |
| | half = int(l // 2) |
| | return s[:half] + s[-half:] if len(s) > l else s |
| |
|
| |
|
| | if __name__ == "__main__": |
| | parser = argparse.ArgumentParser() |
| | parser.add_argument("--input-file", type=str, required=True) |
| | parser.add_argument("--model", type=str, default="gpt-3.5-turbo") |
| | parser.add_argument("--num-prompts", type=int, default=100) |
| | args = parser.parse_args() |
| |
|
| | model = args.model |
| |
|
| | cluster_infos = pickle.load(open(args.input_file, "rb")) |
| | num_total_prompts = sum([x[0] for x in cluster_infos]) |
| |
|
| | topics = [] |
| | percentages = [] |
| | for i, info in enumerate(cluster_infos): |
| | num_samples, topk_prompts, random_prompts = info |
| | percentage = num_samples / num_total_prompts |
| | print( |
| | f"cluster {i}, #prompts {num_samples}, percentage: {percentage * 100:.2f}%" |
| | ) |
| | instruct = "Given a list of user messages, use less than 8 words to summarize a central topic for all messages in English. Your output should only include a single line. Try to be specific." |
| | split = int(args.num_prompts * 0.8) |
| | prompt = "\n".join( |
| | [truncate_string(x, l=200) for x in topk_prompts[:split]] |
| | + [ |
| | truncate_string(x, l=200) |
| | for x in random_prompts[: args.num_prompts - split] |
| | ] |
| | ) |
| | prompt = "BEGIN OF THE MESSAGE LIST\n" + prompt + "\nEND OF THE MESSAGE LIST." |
| |
|
| | if "azure-" in model: |
| | template_name = "chatgpt" |
| | completion_func = chat_completion_openai_azure |
| | elif "gpt" in model: |
| | template_name = "chatgpt" |
| | completion_func = chat_completion_openai |
| | elif "claude" in model: |
| | template_name = "claude" |
| | completion_func = chat_completion_anthropic |
| |
|
| | conv = get_conv_template(template_name) |
| | conv.set_system_message(instruct) |
| | conv.append_message(conv.roles[0], prompt) |
| | conv.append_message(conv.roles[1], None) |
| |
|
| | topic = completion_func(model, conv, temperature=0, max_tokens=256) |
| | print(topic) |
| |
|
| | topics.append(topic) |
| | percentages.append(round(percentage, 6)) |
| |
|
| | print() |
| | print(f"topics: {topics}") |
| | print(f"percentages: {percentages}") |
| |
|
| | |
| | df = pd.DataFrame() |
| | df["topic"] = topics |
| | df["percentage"] = percentages |
| |
|
| | df.to_json(f"cluster_summary_{len(df)}.jsonl", lines=True, orient="records") |
| |
|
