infer_abs_score.py

from model import Model
from data import Preprocess

if __name__ == '__main__':
    checkpoint_path = "pretrained_v3_fp16.pt"
    
    # get model + data processor
    model = Model.load_from_checkpoint(checkpoint_path)
    preprocess = Preprocess(model)
    
    # example input
    tag_str = "safe, rainbow dash, cute, hug, hugging, duo, white background, pony, mare"
    datetime = "2015-10-16T15:51:38Z"
    
    # calculate expected absolute score
    input = preprocess(tag_str, datetime)
    output = model.infer_abs_scores(input)[0]
    print('tags =', tag_str)
    print(f"{output['wilson_score']:.1%} predicted Wilson_Score")
    print(f"{output['score']:.0f} predicted Score")
    print(f"{output['upvotes']:.0f} predicted Upvotes")
    print(f"{output['downvotes']:.0f} predicted Downvotes")
    
    # for this example;
    # >>> tags = safe, rainbow dash, cute, hug, hugging, duo, white background, pony, mare
    # >>> 94.8% predicted Wilson_Score
    # >>> 193 predicted Score
    # >>> 193 predicted Upvotes
    # >>> 1 predicted Downvotes