File size: 2,458 Bytes
82f1bf5
 
 
 
 
 
 
a5d7977
d754e91
62b53be
82f1bf5
 
a5d7977
82f1bf5
 
 
 
 
 
fa0947b
a5d7977
82f1bf5
d754e91
a5d7977
 
82f1bf5
a5d7977
 
 
 
 
 
 
 
 
 
 
 
d754e91
 
82f1bf5
 
 
 
 
 
 
 
90c428d
d754e91
82f1bf5
 
a5d7977
 
 
 
 
 
 
d754e91
82f1bf5
 
 
62b53be
82f1bf5
a5d7977
 
 
d754e91
82f1bf5
 
27c4f46
82f1bf5
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
import os
import sys

import fire
import gradio as gr

from llama_lora.globals import Global
from llama_lora.models import prepare_base_model
from llama_lora.ui.main_page import main_page, get_page_title, main_page_custom_css
from llama_lora.utils.data import init_data_dir



def main(
    base_model: str = "",
    data_dir: str = "",
    # Allows to listen on all interfaces by providing '0.0.0.0'.
    server_name: str = "127.0.0.1",
    share: bool = False,
    skip_loading_base_model: bool = False,
    load_8bit: bool = False,
    ui_show_sys_info: bool = True,
    ui_dev_mode: bool = False,
    wandb_api_key: str = "",
    wandb_project: str = "",
):
    '''
    Start the LLaMA-LoRA Tuner UI.

    :param base_model: (required) The name of the default base model to use.
    :param data_dir: (required) The path to the directory to store data.
    :param server_name: Allows to listen on all interfaces by providing '0.0.0.0'.
    :param share: Create a public Gradio URL.

    :param wandb_api_key: The API key for Weights & Biases. Setting either this or `wandb_project` will enable Weights & Biases.
    :param wandb_project: The default project name for Weights & Biases. Setting either this or `wandb_api_key` will enable Weights & Biases.
    '''

    base_model = base_model or os.environ.get("LLAMA_LORA_BASE_MODEL", "")
    data_dir = data_dir or os.environ.get("LLAMA_LORA_DATA_DIR", "")
    assert (
        base_model
    ), "Please specify a --base_model, e.g. --base_model='decapoda-research/llama-7b-hf'"

    assert (
        data_dir
    ), "Please specify a --data_dir, e.g. --data_dir='./data'"

    Global.default_base_model_name = base_model
    Global.data_dir = os.path.abspath(data_dir)
    Global.load_8bit = load_8bit

    if len(wandb_api_key) > 0:
        Global.enable_wandb = True
        Global.wandb_api_key = wandb_api_key
    if len(wandb_project) > 0:
        Global.enable_wandb = True
        Global.wandb_project = wandb_project

    Global.ui_dev_mode = ui_dev_mode
    Global.ui_show_sys_info = ui_show_sys_info

    os.makedirs(data_dir, exist_ok=True)
    init_data_dir()

    if (not skip_loading_base_model) and (not ui_dev_mode):
        prepare_base_model(base_model)

    with gr.Blocks(title=get_page_title(), css=main_page_custom_css()) as demo:
        main_page()

    demo.queue(concurrency_count=1).launch(server_name=server_name, share=share)


if __name__ == "__main__":
    fire.Fire(main)