Skip to content

Commit f25cd19

Browse files
Olivia-liumalfet
authored andcommitted
Pass port in from cmd. If not passed, default to 5000 (it's the Flask default) (#302)
* chat in browser * pass port in as an arg * more comment
1 parent a9e680d commit f25cd19

File tree

2 files changed

+26
-4
lines changed

2 files changed

+26
-4
lines changed

cli.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99

1010
import torch
1111

12-
default_device = "cpu"
12+
default_device = "cpu"
1313

1414
def check_args(args, name: str) -> None:
1515
pass
@@ -29,8 +29,14 @@ def add_arguments_for_export(parser):
2929
_add_arguments_common(parser)
3030

3131
def add_arguments_for_browser(parser):
32-
# Only export specific options should be here
32+
# Only browser specific options should be here
3333
_add_arguments_common(parser)
34+
parser.add_argument(
35+
"--port",
36+
type=int,
37+
default=5000,
38+
help="Port for the web server for browser mode."
39+
)
3440

3541
def _add_arguments_common(parser):
3642
# TODO: Refactor this so that only common options are here

torchchat.py

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@
2525
parser = argparse.ArgumentParser(description="Top-level command")
2626
subparsers = parser.add_subparsers(
2727
dest="subcommand",
28-
help="Use `generate`, `eval` or `export` followed by subcommand specific options.",
28+
help="Use `generate`, `eval`, `export` or `browser` followed by subcommand specific options.",
2929
)
3030

3131
parser_generate = subparsers.add_parser("generate")
@@ -63,10 +63,26 @@
6363
elif args.subcommand == "browser":
6464
# TODO: add check_args()
6565

66+
# Look for port from cmd args. Default to 5000 if not found.
67+
# The port args will be passed directly to the Flask app.
68+
port = 5000
69+
i = 2
70+
while i < len(sys.argv):
71+
# Check if the current argument is '--port'
72+
if sys.argv[i] == '--port':
73+
# Check if there's a value immediately following '--port'
74+
if i + 1 < len(sys.argv):
75+
# Extract the value and remove '--port' and the value from sys.argv
76+
port = sys.argv[i + 1]
77+
del sys.argv[i:i+2] # Delete '--port' and the value
78+
break # Exit loop since port is found
79+
else:
80+
i += 1
81+
6682
# Assume the user wants "chat" when entering "browser". TODO: add support for "generate" as well
6783
args_plus_chat = ['"{}"'.format(s) for s in sys.argv[2:]] + ['"--chat"'] + ['"--num-samples"'] + ['"1000000"']
6884
formatted_args = ", ".join(args_plus_chat)
69-
command = ["flask", "--app", "chat_in_browser:create_app(" + formatted_args + ")", "run"]
85+
command = ["flask", "--app", "chat_in_browser:create_app(" + formatted_args + ")", "run", "--port", f"{port}"]
7086
subprocess.run(command)
7187
else:
7288
raise RuntimeError("Must specify valid subcommands: generate, export, eval")

0 commit comments

Comments
 (0)