|
|
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() |
|
|
|
