usage

usage.py

@@ -0,0 +1,47 @@
+import os
+import pickle
+from functools import cache
+from pathlib import Path
+
+import numpy as np
+import polars as pl
+from huggingface_hub import hf_hub_download
+
+from embed import embed
+
+DATA = Path(os.environ["DATA_DIR"])
+
+features = ["content", "meta_title", "meta_description"]
+
+
+@cache
+def get_model():
+    file_name = hf_hub_download("opale-ai/news-classifier", "model/model.pickle")
+    with open(file_name, "rb") as f:
+        return pickle.load(f)
+
+
+def record_get():
+    df = pl.read_csv(DATA / "eval.csv")
+    return {col: val for col, val in zip(df.columns, df.sample().row(0))}
+
+
+def record_embed(rec):
+    embeddings = []
+    for f in features:
+        embeddings.append(embed([rec[f]]))
+    return np.hstack(embeddings)
+
+
+def main():
+    model = get_model()
+    record = record_get()
+    embeds = record_embed(record)
+    (pred,) = model.predict(embeds)
+    print(record["content"])
+    print(f"is news (real): {record['is_news_article']}")
+    print(f"is news (pred): {pred}")
+
+
+if __name__ == "__main__":
+    main()