app.py

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 numpy as np
import torch
import wandb
import copy
import os
import sys
import re
from peft import LoraConfig, TaskType, get_peft_model, AutoPeftModelForCausalLM
from sklearn.model_selection import train_test_split
import nltk
from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction

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": "system",
                        "content": """Given the following question-answer pairs, generate 20 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"
                          }
                        ]
                        """
                    },
                    {
                        "role": "user",
                        "content": f"""Here are the question-answer pairs: {json.dumps(self.finetuning_dataset)}
                        """
                    }
                ]
            )
            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):
              row["question"] = row["question"].replace("### Question:", "").replace("### Answer:", "").strip()
              row["answer"] = row["answer"].replace("### Answer:", "").strip()
              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)
            # Create labels (same as input_ids, but mask the non-answer part)
            labels = copy.deepcopy(encoding["input_ids"])
            for i in range(len(examples["question"])):
                # print(examples["question"][i])
                question_length = len(self.tokenizer(examples['question'][i], add_special_tokens=False)["input_ids"])
                # print(f'question length: {question_length}')
                labels[i][:question_length] = [-100] * question_length  # Mask question tokens
            encoding["labels"] = labels
            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
        self.predictions = predictions
        self.labels = labels

        # Convert logits to token IDs using argmax
        predictions = np.argmax(predictions, axis=-1)  

        # Ensure predictions and labels are integers within vocab range
        predictions = np.clip(predictions, 0, self.tokenizer.vocab_size - 1).astype(int)
        labels = np.clip(labels, 0, self.tokenizer.vocab_size - 1).astype(int)

        scores = []

        for prediction, label in zip(predictions, labels):
            print(f"Prediction: {prediction}, Label: {label}")

            # Remove leading 0's from array
            prediction = prediction[np.argmax(prediction != 0):] 
            label = label[np.argmax(label != 0):]

            # Decode predicted tokens
            decoded_preds = self.tokenizer.decode(prediction, skip_special_tokens=True).split()
            decoded_labels = self.tokenizer.decode(label, skip_special_tokens=True).split()

            scores.append(sentence_bleu([decoded_labels], decoded_preds, smoothing_function=SmoothingFunction().method1))

        average_score = sum(scores) / len(scores)
        print(f"Average BLEU score: {average_score}")
        return {"bleu": average_score}

        # return score
        # return {"bleu": 1}

    def train_model(self):
        try:
            tokenized_datasets = self.prepare_data_for_training()
            print('finished preparing 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=5,
                weight_decay=0.01,
                eval_strategy="epoch",
                save_strategy="epoch",
                load_best_model_at_end=True,
            )

            print('training arguments set...')

            # 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],
            )

            print('trainer set...')

            # Start training and yield logs in real-time
            trainer.train()

            # 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 f"✅ Training complete!\n {json.dumps(self.logging_callback.logs)}"
        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("Extend 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. 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("### 6. 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()