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