| import pickle |
| from functools import cache |
|
|
| import polars as pl |
| from huggingface_hub import hf_hub_download |
|
|
| from utils.embed import embed |
| from utils.paths import DATA |
|
|
|
|
| @cache |
| def get_model(): |
| file_name = hf_hub_download( |
| "opale-ai/news-classifier", "model/model.pickle", revision="main" |
| ) |
| with open(file_name, "rb") as f: |
| return pickle.load(f) |
|
|
|
|
| def get_record(): |
| df = pl.read_csv(DATA / "eval.csv") |
| 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"] |
| text = "\n\n".join(rec[k] for k in text_fields) |
| embeds = embed([text]) |
| (pred,) = get_model().predict(embeds) |
| return pred |
|
|
|
|
| def main(): |
| record = get_record() |
| is_news = record["is_news_article"] |
| pred = pred_record(record) |
| print(f"is news (real): {is_news}") |
| print(f"is news (pred): {pred}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
