init

.gitignore

@@ -0,0 +1,4 @@
+__pycache__/
+.venv
+/venv
+.vscode

app.py

@@ -0,0 +1,46 @@
+import os
+import sys
+
+import fire
+import gradio as gr
+
+from llama_lora.globals import Global
+from llama_lora.ui.main_page import main_page
+
+
+def main(
+    load_8bit: bool = False,
+    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,
+    ui_show_sys_info: bool = True,
+):
+    base_model = base_model or os.environ.get("BASE_MODEL", "")
+    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.base_model = base_model
+    Global.data_dir = data_dir
+    Global.data_dir = data_dir
+    Global.load_8bit = load_8bit
+
+    Global.ui_show_sys_info = ui_show_sys_info
+
+    os.makedirs(data_dir, exist_ok=True)
+    os.makedirs(f"{data_dir}/lora", exist_ok=True)
+
+    with gr.Blocks(title=Global.ui_title) as demo:
+        main_page()
+
+    demo.queue().launch(server_name=server_name, share=share)
+
+
+if __name__ == "__main__":
+    fire.Fire(main)

llama_lora/globals.py

@@ -0,0 +1,14 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+
+class Global:
+    base_model: str = ""
+    data_dir: str = ""
+    load_8bit: bool = False
+
+    loaded_tokenizer: Any = None
+    loaded_base_model: Any = None
+
+    # UI related
+    ui_title: str = "LLaMA-LoRA"
+    ui_show_sys_info: bool = True

llama_lora/ui/inference_tab.py

@@ -0,0 +1,10 @@
+import gradio as gr
+
+from ..globals import Global
+
+
+def inference_tab():
+    with gr.Tab("Inference"):
+        gr.Markdown("""
+            WIP.
+            """)

llama_lora/ui/main_page.py

@@ -0,0 +1,21 @@
+import gradio as gr
+
+from ..globals import Global
+
+from .inference_tab import inference_tab
+
+
+def main_page():
+    with gr.Blocks(
+            title="LLaMA-LoRA",
+            css="") as demo:
+        gr.Markdown(f"""
+            # {Global.ui_title}
+
+            Hello world!
+            """)
+        inference_tab()
+        if Global.ui_show_sys_info:
+            gr.Markdown(f"""
+                <small>Data dir: `{Global.data_dir}`</small>
+                """)

requirements.txt

@@ -0,0 +1,12 @@
+accelerate
+appdirs
+bitsandbytes
+black
+black[jupyter]
+datasets
+fire
+git+https://github.com/huggingface/peft.git
+git+https://github.com/huggingface/transformers.git
+gradio
+loralib
+sentencepiece