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