Upload 2 files

app.py

@@ -1,64 +1,137 @@
+import spaces
 import gradio as gr
-from huggingface_hub import InferenceClient
+from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import DataCollatorForSeq2Seq
+from datasets import load_dataset, concatenate_datasets, load_from_disk
+import traceback
 
-"""
-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("HuggingFaceH4/zephyr-7b-beta")
 
+import os
 
-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]})
+@spaces.GPU
+def fine_tune_model(model_name, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
+    try:
+        #login(api_key.strip())
+        # Load the model and tokenizer
+        model = AutoModelForSeq2SeqLM.from_pretrained('google/t5-efficient-tiny-nh8', num_labels=2)
+        
+    
+        # Set training arguments
+        training_args = TrainingArguments(
+            output_dir='/data/results',
+            eval_strategy="steps",  # Change this to steps
+            save_strategy='steps',
+            learning_rate=lr*0.00001,
+            per_device_train_batch_size=int(batch_size),
+            per_device_eval_batch_size=int(batch_size), 
+            num_train_epochs=int(num_epochs),
+            weight_decay=0.01,
+            gradient_accumulation_steps=int(grad),
+            max_grad_norm = 1.0, 
+            load_best_model_at_end=True,
+            metric_for_best_model="accuracy",
+            greater_is_better=True,
+            logging_dir='/data/logs',
+            logging_steps=10,
+            #push_to_hub=True,
+            hub_model_id=hub_id.strip(),
+            fp16=True,
+            #lr_scheduler_type='cosine',
+            save_steps=100,  # Save checkpoint every 500 steps
+            save_total_limit=3, 
+        )
+        # Check if a checkpoint exists and load it
+        max_length = 128
+        # Load the dataset
+        dataset = load_dataset(dataset_name.strip())
+        tokenizer = AutoTokenizer.from_pretrained('google/t5-efficient-tiny-nh8')
+        # Tokenize the dataset
+        def tokenize_function(examples):
+            
+            # Assuming 'text' is the input and 'target' is the expected output
+            model_inputs = tokenizer(
+                examples['text'], 
+                max_length=max_length,  # Set to None for dynamic padding
+                padding=True,     # Disable padding here, we will handle it later
+                truncation=True,
+            )
+        
+            # Setup the decoder input IDs (shifted right)
+            labels = tokenizer(
+                examples['target'], 
+                max_length=max_length,  # Set to None for dynamic padding
+                padding=True,     # Disable padding here, we will handle it later
+                truncation=True,
+                text_target=examples['target']  # Use text_target for target text
+            )
+        
+            # Add labels to the model inputs
+            model_inputs["labels"] = labels["input_ids"]
+        
+        
+        tokenized_datasets = dataset.map(tokenize_function, batched=True)
+        
+        tokenized_datasets['train'].save_to_disk(f'/data/{hub_id.strip()}_train_dataset')
+        tokenized_datasets['test'].save_to_disk(f'/data/{hub_id.strip()}_test_dataset')
+    
+        # Create Trainer
+        trainer = Trainer(
+            model=model,
+            args=training_args,
+            train_dataset=tokenized_datasets['train'],
+            eval_dataset=tokenized_datasets['test'],
+            compute_metrics=compute_metrics,
+            #callbacks=[LoggingCallback()], 
+        )            
 
-    messages.append({"role": "user", "content": message})
+        # Fine-tune the model
+        trainer.train()
+        trainer.push_to_hub(commit_message="Training complete!")
+    except Exception as e:
+        return f"An error occurred: {str(e)}, TB: {traceback.format_exc()}"
+    return 'DONE!'#model
+'''
+# Define Gradio interface
+def predict(text):
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name.strip(), num_labels=2)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    outputs = model(inputs)
+    predictions = outputs.logits.argmax(dim=-1)
+    return predictions.item()
+'''
+# Create Gradio interface
+try:
+    iface = gr.Interface(
+        fn=fine_tune_model,
+        inputs=[
+            gr.Textbox(label="Model Name (e.g., 'google/t5-efficient-tiny-nh8')"),
+            gr.Textbox(label="Dataset Name (e.g., 'imdb')"),
+            gr.Textbox(label="HF hub to push to after training"),
+            gr.Textbox(label="HF API token"),
+            gr.Slider(minimum=1, maximum=10, value=3, label="Number of Epochs", step=1),
+            gr.Slider(minimum=1, maximum=2000, value=1, label="Batch Size", step=1),
+            gr.Slider(minimum=1, maximum=1000, value=1, label="Learning Rate (e-5)", step=1),
+            gr.Slider(minimum=1, maximum=100, value=1, label="Gradient accumulation", step=1), 
+        ],
+        outputs="text",
+        title="Fine-Tune Hugging Face Model",
+        description="This interface allows you to fine-tune a Hugging Face model on a specified dataset."
+    )
+    '''
+    iface = gr.Interface(
+        fn=predict,
+        inputs=[
+            gr.Textbox(label="Query"),
+        ],
+        outputs="text",
+        title="Fine-Tune Hugging Face Model",
+        description="This interface allows you to test a fine-tune Hugging Face model."
+    )
+    '''
+    # Launch the interface
+    iface.launch()    
+except Exception as e:
+    print(f"An error occurred: {str(e)}, TB: {traceback.format_exc()}")
 
-    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
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = 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)",
-        ),
-    ],
-)
-
-
-if __name__ == "__main__":
-    demo.launch()

requirements.txt

@@ -1 +1,8 @@
-huggingface_hub==0.25.2
+spaces
+transformers
+datasets
+peft
+huggingface_hub
+scikit-learn
+numpy
+torch