File size: 3,197 Bytes
df1ce08
 
8245e16
 
df1ce08
 
 
 
f934c1a
df1ce08
19f8b29
df1ce08
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19f8b29
f934c1a
 
9ad7da7
 
8245e16
f934c1a
 
8245e16
 
f934c1a
8245e16
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f934c1a
8245e16
 
 
 
 
 
19f8b29
df1ce08
 
 
8245e16
 
 
 
 
 
 
 
19f8b29
 
 
 
 
 
f934c1a
 
 
 
 
 
 
8245e16
 
 
9ad7da7
8245e16
 
df1ce08
f934c1a
df1ce08
 
8245e16
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
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
import gradio as gr
from huggingface_hub import InferenceClient
from transformers import AutoModelForCausalLM, Trainer, TrainingArguments
from datasets import load_dataset

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
client = InferenceClient("meta-llama/Meta-Llama-3-8B-Instruct")


def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content

        response += token
        yield response


def train_model():

    os.environ["HUGGINGFACE_TOKEN"] = hf_token

    # Load dataset
    dataset = load_dataset('json', data_files={
                           'train': 'training_set.json'})

    # Load model
    model = AutoModelForCausalLM.from_pretrained('meta-llama/Meta-Llama-3-8B-Instruct')

    # Define training arguments
    training_args = TrainingArguments(
        output_dir='./results',
        num_train_epochs=3,
        per_device_train_batch_size=16,
        save_steps=10_000,
        save_total_limit=2,
    )

    # Initialize Trainer
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=dataset['train'],
        eval_dataset=dataset['test']
    )

    # Start training
    trainer.train()
    return "Training complete"


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.Blocks()

with demo:
    gr.Markdown("# Llama3training Chatbot and Model Trainer")
    with gr.Tab("Chat"):
        gr.ChatInterface(
            respond,
            additional_inputs=[
                gr.Textbox(value="You are a friendly Chatbot.",
                           label="System message"),
                gr.Slider(minimum=1, maximum=2048, value=512,
                          step=1, label="Max new tokens"),
                gr.Slider(minimum=0.1, maximum=4.0, value=0.7,
                          step=0.1, label="Temperature"),
                gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.95,
                    step=0.05,
                    label="Top-p (nucleus sampling)",
                ),
            ],
        )
    with gr.Tab("Train"):
        hf_token = gr.Textbox(label="Hugging Face Token", type="password")
        train_button = gr.Button("Start Training")
        train_output = gr.Textbox(label="Training Output")

        train_button.click(train_model, outputs=train_output)

if __name__ == "__main__":
    demo.launch()