Add path argument (#888)

Co-authored-by: Abhinav Rao <t-raoabhinav@microsoft.com>
alpa-projects · Feb 28, 2023 · 06cb62f · 06cb62f
1 parent f76fe8d
commit 06cb62f
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/examples/llm_serving/textgen.py b/examples/llm_serving/textgen.py
@@ -26,7 +26,7 @@ def main(args):
 
     # Load the model
     model = get_model(model_name=args.model,
-                      path="~/opt_weights",
+                      path=args.path,
                       batch_size=args.n_prompts,
                       **generate_params)
 
@@ -53,7 +53,8 @@ def main(args):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model", type=str, default="alpa/opt-1.3b")
+    parser.add_argument('--model', type=str, default='alpa/opt-1.3b')
+    parser.add_argument('--path', type=str, default='~/opt_weights')
     parser.add_argument('--do-sample', action='store_true')
     parser.add_argument('--num-beams', type=int, default=1)
     parser.add_argument('--num-return-sequences', type=int, default=1)