usage

usage.py

@@ -19,20 +19,22 @@ def get_model():
 
 def get_record():
     df = pl.read_csv(DATA / "eval.csv")
-    raw = {col: val for col, val in zip(df.columns, df.sample().row(0))}
+    return {col: val for col, val in zip(df.columns, df.sample().row(0))}
+
+
+def pred_record(rec):
     text_fields = ["meta_title", "meta_description", "content"]
-    return {
-        "text": "\n\n".join(raw[k] for k in text_fields),
-        "is_news": raw["is_news_article"],
-    }
+    text = "\n\n".join(rec[k] for k in text_fields)
+    embeds = embed([text])
+    (pred,) = get_model().predict(embeds)
+    return pred
 
 
 def main():
-    model = get_model()
     record = get_record()
-    embeds = embed([record["text"]])
-    (pred,) = model.predict(embeds)
-    print(f"is news (real): {record['is_news']}")
+    is_news = record["is_news_article"]
+    pred = pred_record(record)
+    print(f"is news (real): {is_news}")
     print(f"is news (pred): {pred}")