import argparse
import torch
from sentence_transformers import CrossEncoder
from openmind import is_torch_npu_available

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to model",
        default=None,
    )
    args = parser.parse_args()
    return args

def main():
    args = parse_args()
    MODEL_NAME = args.model_name_or_path
    if is_torch_npu_available():
        device = "npu:0"
    else:
        device = "cpu"
    model = CrossEncoder(MODEL_NAME, max_length=512, device=device)
    model.model.half()
    query = "感動的な映画について"
    passages = [
        "深いテーマを持ちながらも、観る人の心を揺さぶる名作。登場人物の心情描写が秀逸で、ラストは涙なしでは見られない。",
        "重要なメッセージ性は評価できるが、暗い話が続くので気分が落ち込んでしまった。もう少し明るい要素があればよかった。",
        "どうにもリアリティに欠ける展開が気になった。もっと深みのある人間ドラマが見たかった。",
        "アクションシーンが楽しすぎる。見ていて飽きない。ストーリーはシンプルだが、それが逆に良い。",
    ]
    scores = model.predict([(query, passage) for passage in passages])
    print(scores)
    
if __name__ == "__main__":
    main()