app code change for train token

app.py

@@ -2,11 +2,12 @@ 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")
+client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
 
 
 def respond(
@@ -42,13 +43,16 @@ def respond(
         yield response
 
 
-def train_model():
+def train_model(hf_token):
+    # Set the Hugging Face token as an environment variable
+    os.environ["HUGGINGFACE_TOKEN"] = hf_token
+
     # Load dataset
-    dataset = load_dataset('json', data_files={
-                           'train': 'training_set.json'})
+    dataset = load_dataset('json', data_files='dataset.jsonl')
 
     # Load model
-    model = AutoModelForCausalLM.from_pretrained('meta-llama/Meta-Llama-3-8B-Instruct')
+    model = AutoModelForCausalLM.from_pretrained(
+        'meta-llama/Meta-Llama-3-8B-Instruct', use_auth_token=hf_token)
 
     # Define training arguments
     training_args = TrainingArguments(
@@ -64,7 +68,8 @@ def train_model():
         model=model,
         args=training_args,
         train_dataset=dataset['train'],
-        eval_dataset=dataset['test']
+        # Using train as eval for this simple example
+        eval_dataset=dataset['train']
     )
 
     # Start training
@@ -89,20 +94,17 @@ with demo:
                           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)",
-                ),
+                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)
+        train_button.click(fn=train_model, inputs=hf_token,
+                           outputs=train_output)
 
 if __name__ == "__main__":
     demo.launch()