import argparse
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM

def main():
    parser = argparse.ArgumentParser(description="Run SASOK CLI")
    parser.add_argument("prompt", type=str, help="Input prompt for generation")
    parser.add_argument("--model", type=str, default="sasok_model", help="Model directory or HuggingFace model ID")
    args = parser.parse_args()

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForCausalLM.from_pretrained(args.model)

    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
    result = pipe(args.prompt, max_length=100, do_sample=True)
    print(result[0]["generated_text"])

if __name__ == "__main__":
    main()