File size: 3,252 Bytes
df1ce08
 
8245e16
 
0925243
df1ce08
 
 
 
f934c1a
df1ce08
19f8b29
df1ce08
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
19f8b29
a4a6c55
 
9ad7da7
8245e16
f934c1a
 
8245e16
 
846bac8
 
8245e16
 
 
 
 
 
 
 
 
 
 
 
 
 
 
f934c1a
8245e16
 
 
 
 
 
19f8b29
df1ce08
 
 
8245e16
 
 
 
 
 
 
 
19f8b29
 
 
 
 
 
f934c1a
 
 
 
 
 
 
8245e16
 
 
9ad7da7
8245e16
 
df1ce08
a4a6c55
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
114
import gradio as gr
from huggingface_hub import InferenceClient
from transformers import AutoModelForCausalLM, Trainer, TrainingArguments
from datasets import load_dataset
import os

"""
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(hf_token_value):
    os.environ["HUGGINGFACE_TOKEN"] = hf_token_value

    # 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, inputs=hf_token, outputs=train_output)

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