feat: add eval script

Former-commit-id: 9ca03cf7fd6b5cebbb24e023ab238eafd11ed76e [formerly e2d22f446a0f4f0a2817dadb84684eec596c5bc6]
Former-commit-id: 8362e9b7fd02c0f7ce0d4f350e7b7758c97952ef

planning_ai/eval/compare_summaries.py

@@ -6,53 +6,69 @@ from pydantic import BaseModel, Field
 from planning_ai.common.utils import Paths
 from planning_ai.llms.llm import GPT4o
 
-with open("./planning_ai/eval/eval.txt", "r") as f:
-    compare_template = f.read()
-
-with open("./planning_ai/eval/summary.txt", "r") as f:
-    summary_template = f.read()
-
 
 class SummaryEvaluator(BaseModel):
     score: int = Field(..., description="The number of the best summary.")
 
 
-SLLM = GPT4o.with_structured_output(SummaryEvaluator, strict=True)
+def load_templates():
+    with open("./planning_ai/eval/eval.txt", "r") as f:
+        compare_template = f.read()
+    with open("./planning_ai/eval/summary.txt", "r") as f:
+        summary_template = f.read()
+    return compare_template, summary_template
+
 
-compare_prompt = ChatPromptTemplate([("system", compare_template)])
-compare_chain = compare_prompt | SLLM
+def initialize_chains(compare_template, summary_template):
+    SLLM = GPT4o.with_structured_output(SummaryEvaluator, strict=True)
+    compare_prompt = ChatPromptTemplate([("system", compare_template)])
+    compare_chain = compare_prompt | SLLM
 
-summary_prompt = ChatPromptTemplate([("system", summary_template)])
-summary_chain = summary_prompt | GPT4o | StrOutputParser()
+    summary_prompt = ChatPromptTemplate([("system", summary_template)])
+    summary_chain = summary_prompt | GPT4o | StrOutputParser()
 
+    return compare_chain, summary_chain
 
-original = pl.read_parquet(Paths.STAGING / "gcpt3.parquet").filter(
-    pl.col("attachments_id").is_null()
-)
-summaries1 = original[["text", "representations_summary"]].unique().head(20)
 
-summaries2 = summaries1[["text"]]
-summaries2 = summaries2.with_columns(
-    pl.col("text")
-    .map_elements(
-        lambda x: summary_chain.invoke({"content": x}), return_dtype=pl.String
+def process_summaries(compare_chain, summary_chain):
+    original = pl.read_parquet(Paths.STAGING / "gcpt3.parquet").filter(
+        pl.col("attachments_id").is_null()
+    )
+    summaries1 = original[["text", "representations_summary"]].unique()
+
+    summaries2 = summaries1[["text"]]
+    summaries2 = summaries2.with_columns(
+        pl.col("text")
+        .map_elements(
+            lambda x: summary_chain.invoke({"content": x}), return_dtype=pl.String
+        )
+        .alias("summary")
     )
-    .alias("summary")
-)
-
-summaries = summaries1.join(summaries2, on="text")
-summaries = summaries.with_columns(
-    pl.struct(["text", "representations_summary", "summary"])
-    .map_elements(
-        lambda x: compare_chain.invoke(
-            {
-                "document": x["text"],
-                "summary_1": x["representations_summary"],
-                "summary_2": x["summary"],
-            }
-        ).score,
-        return_dtype=pl.Int8,
+
+    summaries = summaries1.join(summaries2, on="text")
+    summaries = summaries.with_columns(
+        pl.struct(["text", "representations_summary", "summary"])
+        .map_elements(
+            lambda x: compare_chain.invoke(
+                {
+                    "document": x["text"],
+                    "summary_1": x["representations_summary"],
+                    "summary_2": x["summary"],
+                }
+            ).score,
+            return_dtype=pl.Int8,
+        )
+        .alias("score")
     )
-    .alias("score")
-)
-summaries["score"].value_counts()
+    return summaries
+
+
+def main():
+    compare_template, summary_template = load_templates()
+    compare_chain, summary_chain = initialize_chains(compare_template, summary_template)
+    summaries = process_summaries(compare_chain, summary_chain)
+    summaries.write_parquet(Paths.OUT / "eval.parquet")
+
+
+if __name__ == "__main__":
+    main()