diff --git a/optillm.py b/optillm.py index 005e866f..e19a5960 100644 --- a/optillm.py +++ b/optillm.py @@ -649,7 +649,8 @@ def parse_args(): ("--n", "OPTILLM_N", int, 1, "Number of final responses to be returned"), ("--return-full-response", "OPTILLM_RETURN_FULL_RESPONSE", bool, False, "Return the full response including the CoT with tags"), ("--port", "OPTILLM_PORT", int, 8000, "Specify the port to run the proxy"), - ("--log", "OPTILLM_LOG", str, "info", "Specify the logging level", list(logging_levels.keys())) + ("--log", "OPTILLM_LOG", str, "info", "Specify the logging level", list(logging_levels.keys())), + ("--launch-gui", "OPTILLM_LAUNCH_GUI", bool, False, "Launch a Gradio chat interface") ] for arg, env, type_, default, help_text, *extra in args_env: @@ -688,8 +689,8 @@ def parse_args(): def main(): global server_config # Call this function at the start of main() - load_plugins() args = parse_args() + load_plugins() # Update server_config with all argument values server_config.update(vars(args)) @@ -706,6 +707,32 @@ def main(): if server_config_clean['optillm_api_key']: server_config_clean['optillm_api_key'] = '[REDACTED]' logger.info(f"Server configuration: {server_config_clean}") + + # Launch GUI if requested + if server_config.get('launch_gui'): + try: + import gradio as gr + # Start server in a separate thread + import threading + server_thread = threading.Thread(target=app.run, kwargs={'host': '0.0.0.0', 'port': port}) + server_thread.daemon = True + server_thread.start() + + # Configure the base URL for the Gradio interface + base_url = f"http://localhost:{port}/v1" + logger.info(f"Launching Gradio interface connected to {base_url}") + + # Launch Gradio interface + demo = gr.load_chat( + base_url, + model=server_config['model'], + token=None + ) + demo.launch(server_name="0.0.0.0", share=False) + except ImportError: + logger.error("Gradio is required for GUI. Install it with: pip install gradio") + return + app.run(host='0.0.0.0', port=port) if __name__ == "__main__": diff --git a/requirements.txt b/requirements.txt index 1d710c6f..af328f48 100644 --- a/requirements.txt +++ b/requirements.txt @@ -20,4 +20,5 @@ nbconvert ipython ipykernel peft -bitsandbytes \ No newline at end of file +bitsandbytes +gradio \ No newline at end of file diff --git a/setup.py b/setup.py index cb6e81da..db3b1240 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ setup( name="optillm", - version="0.0.21", + version="0.0.22", packages=find_packages(), py_modules=['optillm'], package_data={ @@ -33,6 +33,7 @@ "ipykernel", "peft", "bitsandbytes", + "gradio", ], entry_points={ 'console_scripts': [