cold/__main__.py

import argparse
from .classifier import ToxicTextClassifier

def getArgs():
    parser = argparse.ArgumentParser(description="LiteDetective - Malicious Content Detection Pipeline")
    parser.add_argument("--path", type=str, default="output/cold.pth", required=False, help="Path to the model")
    parser.add_argument("--device", type=str, default="cpu", required=False, help="Device to use (cpu, mps, or cuda)")
    parser.add_argument("args", nargs='+', help="the text to detect")

    args = parser.parse_args()
    return args

def main():
    args = getArgs()
    model = ToxicTextClassifier(path=args.path)
    result = model.predict(args.args, device=args.device)
    print(result)

if __name__ == "__main__":
    main()