Create app.py

app.py

@@ -0,0 +1,329 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, TrainerCallback, DataCollatorWithPadding, DefaultDataCollator
+from openai import OpenAI
+from huggingface_hub import login
+import datasets
+from datasets import Dataset
+import json
+import pandas as pd
+import torch
+import wandb
+import os
+import sys
+from peft import LoraConfig, TaskType, get_peft_model, AutoPeftModelForCausalLM
+from sklearn.model_selection import train_test_split
+
+IS_COLAB = False
+if "google.colab" in sys.modules or "google.colab" in os.environ:
+    IS_COLAB = True
+
+# Load env secrets
+if IS_COLAB:
+    from google.colab import userdata
+    OPENAI_API_KEY=userdata.get('OPENAI_API_KEY')
+    WANDB_API_KEY=userdata.get('WANDB_API_KEY')
+else:
+    OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
+    WANDB_API_KEY = os.environ.get("WANDB_API_KEY")
+
+# Authenticate Weights and Biases
+wandb.login(key=WANDB_API_KEY)
+
+# Custom callback to capture logs
+class LoggingCallback(TrainerCallback):
+    def __init__(self):
+        self.logs = []  # Store logs
+
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if logs:
+            self.logs.append(logs)  # Append logs to list
+
+
+class LLMTrainingApp:
+    def __init__(self):
+        # self.metric = datasets.load_metric('sacrebleu')
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.finetuning_dataset = []
+        self.prompt_template = """### Question: {question} ### Answer: """
+        self.training_output = "/content/peft-model" if IS_COLAB else "./peft-model"
+        self.localpath = "/content/finetuned-model" if IS_COLAB else "./finetuned-model"
+        self.tokenizer = None
+        self.model = None
+        self.model_name = None
+        self.fine_tuned_model = None
+        self.teacher_model = OpenAI(api_key=OPENAI_API_KEY)
+        self.base_models = {
+                  "SmolLM": {"hf_name":"HuggingFaceTB/SmolLM2-135M",
+                              "model_size": "135M",
+                              "training_size": "2T",
+                              "context_window": "8192"},
+                  "GPT2": {"hf_name":"openai-community/gpt2",
+                              "model_size": "137M",
+                              "training_size": "2T",
+                              "context_window": "1024"}
+                  }
+        self.peft_config = LoraConfig(task_type=TaskType.CAUSAL_LM, inference_mode=False, r=8, lora_alpha=32, lora_dropout=0.1)
+        self.logging_callback = LoggingCallback()
+
+    def login_into_hf(self, token):
+          if not token:
+                return "❌ Please enter a valid token."
+          try:
+              login(token)
+              return f"✅ Logged in successfully!"
+          except Exception as e:
+              return f"❌ Login failed: {str(e)}"
+
+    def select_model(self, model_name):
+        self.model_name = model_name
+        model_hf_name = self.base_models[model_name]["hf_name"]
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(model_hf_name)
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            base_model = AutoModelForCausalLM.from_pretrained(
+                model_hf_name,
+                torch_dtype="auto",
+                device_map="auto"
+            )
+            self.model = get_peft_model(base_model, self.peft_config)
+            params = self.model.get_nb_trainable_parameters()
+            percent_trainable = round(100 * (params[0] / params[1]), 2)
+            return f"✅ Loaded model into memory! Base Model card: {json.dumps(self.base_models[model_name])} - % of trainable parameters for PEFT model: {percent_trainable}"
+        except Exception as e:
+              return f"❌ Failed to load model and/or tokenizer: {str(e)}"
+
+    def create_golden_dataset(self, dataset):
+        try:
+            dataset = pd.DataFrame(dataset)
+            for i, row in dataset.iterrows():
+              self.finetuning_dataset.append({"question": self.prompt_template.format(question=row["Question"]), "answer": row["Answer"]})
+            return "✅ Golden dataset created!"
+        except Exception as e:
+              return f"❌ Failed to create dataset: {str(e)}"
+
+    def extend_dataset(self):
+        try:
+            completion = self.teacher_model.chat.completions.create(
+                model="gpt-4o",
+                messages=[
+                    {
+                        "role": "user",
+                        "content": """Given the following question-answer pairs, generate 10 similar pairs in the following json format below. Do not respond with anything other than the json.
+                        ```json
+                        [
+                          {
+                            "question": "question 1",
+                            "answer": "answer 1"
+                          },
+                          {
+                            "question": "question 2",
+                            "answer": "answer 2"
+                          }
+                        ]
+                        """
+                    }
+                ]
+            )
+            response = completion.choices[0].message.content
+            print(f"raw response: {response}")
+            clean_response = response.replace("```json", "").replace("```", "").strip()
+            print(f"clean response: {clean_response}")
+            new_data = json.loads(clean_response)
+            for i, row in enumerate(new_data):
+              self.finetuning_dataset.append({"question": self.prompt_template.format(question=row["question"]), "answer": row["answer"]})
+            # create df to display
+            df = pd.DataFrame(new_data)
+            return "✅ Synthetic dataset generated!", df
+        except Exception as e:
+            return f"❌ Failed to generate synthetic dataset: {str(e)}", pd.DataFrame()
+
+    def tokenize_function(self, examples):
+        try:
+            # Tokenize the question and answer as input and target (labels) for causal LM
+            encoding = self.tokenizer(examples['question'], examples['answer'], padding=True)
+            # Set the labels as the input_ids
+            encoding['labels'] = encoding['input_ids'].copy()
+            return encoding
+        except Exception as e:
+            return f"❌ Failed to tokenize input: {str(e)}"
+
+
+    def prepare_data_for_training(self):
+        try:
+            dataset = Dataset.from_dict({
+            "question": [entry["question"] for entry in self.finetuning_dataset],
+            "answer": [entry["answer"] for entry in self.finetuning_dataset],
+            })
+            dataset = dataset.map(self.tokenize_function, batched=True)
+            train_dataset, test_dataset = dataset.train_test_split(test_size=0.2).values()
+            return {"train": train_dataset, "test": test_dataset}
+        except Exception as e:
+            return f"❌ Failed to prepare data for training: {str(e)}"
+
+
+    def compute_bleu(self, eval_pred):
+        predictions, labels = eval_pred
+        # # Flatten predictions and labels if they are in nested lists
+        # predictions = predictions.flatten()
+        # labels = labels.flatten()
+        # # Ensure that predictions and labels are integers
+        # predictions = predictions.astype(int)  # Convert to integer
+        # labels = labels.astype(int)  # Convert to integer
+        # # Decode the predicted tokens
+        # decoded_preds = self.tokenizer.decode(predictions, skip_special_tokens=True)
+        # decoded_labels = self.tokenizer.decode(labels, skip_special_tokens=True)
+        # result = self.metric.compute(predictions=[decoded_preds], references=[[decoded_labels]])
+        result = {"bleu": 1}
+        return result
+
+    def log_generator(self):
+        """ Continuously send logs to frontend during training """
+        for log in self.logging_callback.logs:
+            yield str(log)
+
+    def train_model(self):
+        try:
+            tokenized_datasets = self.prepare_data_for_training()
+
+            # Create training arguments
+            training_args = TrainingArguments(
+                output_dir=self.training_output,
+                learning_rate=1e-3,
+                per_device_train_batch_size=32,
+                per_device_eval_batch_size=32,
+                num_train_epochs=2,
+                weight_decay=0.01,
+                eval_strategy="epoch",
+                save_strategy="epoch",
+                load_best_model_at_end=True,
+            )
+
+            # Create trainer & attach logging callback
+            trainer = Trainer(
+                model=self.model,
+                args=training_args,
+                train_dataset=tokenized_datasets["train"],
+                eval_dataset=tokenized_datasets["test"],
+                tokenizer=self.tokenizer,
+                data_collator=DefaultDataCollator(),
+                compute_metrics=self.compute_bleu,
+                callbacks=[self.logging_callback],
+            )
+
+            # Start training and yield logs in real-time
+            trainer.train()
+
+            # for log in logging_callback.logs:
+            #     yield str(log)
+
+            # Save trained model to HF
+            self.model.save_pretrained(self.localpath) # save to local
+            self.model.push_to_hub(f"{self.model_name}-lora")
+
+            return "✅ Training complete!"
+        except Exception as e:
+            return f"❌ Training failed: {str(e)}"
+
+    def run_inference(self, prompt):
+        try:
+            # Load fine-tuned memory into memory and set mode to eval
+            self.fine_tuned_model = AutoPeftModelForCausalLM.from_pretrained(self.localpath)
+            self.fine_tuned_model = self.fine_tuned_model.to(self.device)
+            self.fine_tuned_model.eval()
+
+            # Tokenize input with padding and attention mask
+            inputs = self.tokenizer(prompt, return_tensors="pt", padding=True).to(self.device)
+
+            # Generate response
+            output = self.fine_tuned_model.generate(
+                **inputs,
+                max_length=50,  # Limit response length
+                num_return_sequences=1,  # Single response
+                temperature=0.7,  # Sampling randomness
+                top_p=0.9  # Nucleus sampling
+            )
+
+            response = self.tokenizer.batch_decode(output.detach().cpu().numpy(), skip_special_tokens=True)[0]
+            return response
+        except Exception as e:
+            return f"❌ Inference failed: {str(e)}"
+
+    def build_ui(self):
+        with gr.Blocks() as demo:
+            gr.Markdown("# LLM Fine-tuning")
+
+            # Model Selection
+            with gr.Group():
+                gr.Markdown("### 1. Login into Hugging Face")
+                with gr.Column():
+                    token = gr.Textbox(label="Enter Hugging Face Access Token (w/ write permissions)", type="password")
+                    inference_btn = gr.Button("Login", variant="primary")
+                    status = gr.Textbox(label="Status")
+                inference_btn.click(self.login_into_hf, inputs=token, outputs=status)
+
+            # Model Selection
+            with gr.Group():
+                gr.Markdown("### 2. Select Model")
+                with gr.Column():
+                    model_dropdown = gr.Dropdown([key for key in self.base_models.keys()], label="Small Models")
+                    select_model_btn = gr.Button("Select", variant="primary")
+                    selected_model_text = gr.Textbox(label="Model Status")
+                select_model_btn.click(self.select_model, inputs=model_dropdown, outputs=[selected_model_text])
+
+            # Create Golden Dataset
+            with gr.Group():
+                gr.Markdown("### 3. Create Golden Dataset")
+                with gr.Column():
+                    dataset_table = gr.Dataframe(
+                        headers=["Question", "Answer"],
+                        value=[["", ""] for _ in range(3)],
+                        label="Golden Dataset"
+                    )
+                    create_data_btn = gr.Button("Create Dataset", variant="primary")
+                    dataset_status = gr.Textbox(label="Dataset Status")
+                create_data_btn.click(self.create_golden_dataset, inputs=dataset_table, outputs=[dataset_status])
+
+            # Generate Full Dataset
+            with gr.Group():
+                gr.Markdown("### 4. Extend Dataset with Synthetic Data")
+                with gr.Column():
+                    dataset_table = gr.Dataframe(
+                        headers=["Question", "Answer"],
+                        label="Golden + Synthetic Dataset"
+                    )
+                    generate_status = gr.Textbox(label="Dataset Generation Status")
+                    generate_data_btn = gr.Button("Generate Dataset", variant="primary")
+                generate_data_btn.click(self.extend_dataset, outputs=[generate_status, dataset_table])
+
+            # Train Model & Visualize Loss
+            with gr.Group():
+                gr.Markdown("### 5. Start Logging")
+                with gr.Column():
+                    train_status = gr.Textbox(label="Training Status", lines=10)
+                    train_btn = gr.Button("Train", variant="primary")
+                train_btn.click(self.log_generator, outputs=[train_status])
+
+            # Train Model & Visualize Loss
+            with gr.Group():
+                gr.Markdown("### 6. Train Model")
+                with gr.Column():
+                    train_status = gr.Textbox(label="Training Status")
+                    train_btn = gr.Button("Train", variant="primary")
+                train_btn.click(self.train_model, outputs=[train_status])
+
+            # Run Inference
+            with gr.Group():
+                gr.Markdown("### 7. Run Inference")
+                with gr.Column():
+                    user_prompt = gr.Textbox(label="Enter Prompt")
+                    inference_btn = gr.Button("Run Inference", variant="primary")
+                    inference_output = gr.Textbox(label="Inference Output")
+                inference_btn.click(self.run_inference, inputs=user_prompt, outputs=inference_output)
+
+            return demo
+
+# Create an instance of the app
+app = LLMTrainingApp()
+
+# Launch the Gradio app using the class method
+app.build_ui().launch()