app.py

import gradio as gr
import spaces

import os
import gc
import random
import warnings

warnings.filterwarnings("ignore")

import numpy as np
import pandas as pd

pd.set_option("display.max_rows", 500)
pd.set_option("display.max_columns", 500)
pd.set_option("display.width", 1000)
from tqdm.auto import tqdm

import torch
import torch.nn as nn
import tokenizers
import transformers

print(f"tokenizers.__version__: {tokenizers.__version__}")
print(f"transformers.__version__: {transformers.__version__}")
print(f"torch.__version__: {torch.__version__}")
print(f"torch cuda version: {torch.version.cuda}")
from transformers import AutoTokenizer, AutoConfig
from transformers import BitsAndBytesConfig, AutoModelForCausalLM, MistralForCausalLM
from peft import LoraConfig, get_peft_model


title = "H2O AI Predict the LLM"

description =" The objective of this [competition](https://www.kaggle.com/competitions/h2oai-predict-the-llm) was to \
detect which out of 7 possible LLM models produced a particular response. \n\n\
This demo is utilizing finetuned HuggingFaceH4/zephyr-7b-beta model for a multiclass classification task. \n\n \
We ranked 3rd out of more than 100 participants and our team's solution is [here](https://www.kaggle.com/competitions/h2oai-predict-the-llm/discussion/453728)"

title = title + "\n" + description

#Theme from - https://huggingface.co/spaces/trl-lib/stack-llama/blob/main/app.py
theme = gr.themes.Monochrome(
    primary_hue="indigo",
    secondary_hue="blue",
    neutral_hue="slate",
    radius_size=gr.themes.sizes.radius_sm,
    font=[gr.themes.GoogleFont("Open Sans"), "ui-sans-serif", "system-ui", "sans-serif"],
)

### Load the model
class CFG:
    num_workers = os.cpu_count()
    llm_backbone = "HuggingFaceH4/zephyr-7b-beta" #"save_pretrained_model/zephyr-7b-beta"
    tokenizer_path = "HuggingFaceH4/zephyr-7b-beta"
    tokenizer = AutoTokenizer.from_pretrained(
        tokenizer_path, add_prefix_space=False, use_fast=True, trust_remote_code=True, add_eos_token=True
    )
    batch_size = 1
    max_len = 650
    seed = 42

    num_labels = 7

    lora = True
    lora_r = 4
    lora_alpha = 16
    lora_dropout = 0.05
    lora_target_modules = ""
    gradient_checkpointing = True


class CustomModel(nn.Module):
    """
    Model for causal language modeling problem type.
    """

    def __init__(self):
        super().__init__()

        self.backbone_config = AutoConfig.from_pretrained(
            CFG.llm_backbone, trust_remote_code=True
        )

        quantization_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16,
            bnb_4bit_quant_type="nf4",
        )

        self.model = AutoModelForCausalLM.from_pretrained(
            CFG.llm_backbone,
            config=self.backbone_config,
            quantization_config=quantization_config,
        )

        if CFG.lora:
            target_modules = []
            for name, module in self.model.named_modules():
                if (
                    isinstance(module, (torch.nn.Linear, torch.nn.Conv1d))
                    and "head" not in name
                ):
                    name = name.split(".")[-1]
                    if name not in target_modules:
                        target_modules.append(name)

            lora_config = LoraConfig(
                r=CFG.lora_r,
                lora_alpha=CFG.lora_alpha,
                target_modules=target_modules,
                lora_dropout=CFG.lora_dropout,
                bias="none",
                task_type="CAUSAL_LM",
            )
            if CFG.gradient_checkpointing:
                self.model.enable_input_require_grads()
            self.model = get_peft_model(self.model, lora_config)
            self.model.print_trainable_parameters()

        self.classification_head = nn.Linear(
            self.backbone_config.vocab_size, CFG.num_labels, bias=False
        )
        self._init_weights(self.classification_head)
    
    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=self.backbone_config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.backbone_config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

    def forward(
        self,
        batch
    ):
        # disable cache if gradient checkpointing is enabled
        if CFG.gradient_checkpointing:
            self.model.config.use_cache = False

        self.model.config.pretraining_tp = 1

        output = self.model(
            input_ids=batch["input_ids"],
            attention_mask=batch["attention_mask"],
        )

        output.logits = self.classification_head(output[0][:, -1].float())

        # enable cache again if gradient checkpointing is enabled
        if CFG.gradient_checkpointing:
            self.model.config.use_cache = True

        return output.logits
    
model = CustomModel()
### End Load the model

def do_inference(full_text):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model_paths = [
    'model_finetuned/HuggingFaceH4-zephyr-7b-beta_fold0_best.pth']

    def prepare_input(cfg, text):
        inputs = cfg.tokenizer.encode_plus(
            text,
            return_tensors=None,
            add_special_tokens=True,
            max_length=CFG.max_len,
            pad_to_max_length=True,
            truncation="longest_first",
        )
        for k, v in inputs.items():
            inputs[k] = torch.tensor(v, dtype=torch.long)
        return inputs

    state = torch.load(model_paths[0], map_location=torch.device("cpu"))
    model.load_state_dict(state["model"] ,strict=False)
    model.eval()
    model.to(device)

    inputs = prepare_input(CFG, full_text)
    inputs["input_ids"] = inputs["input_ids"].reshape(1, -1).to(device)
    inputs["attention_mask"] = inputs["attention_mask"].reshape(1, -1).to(device)
    
    with torch.no_grad():
        with torch.cuda.amp.autocast(
            enabled=True, dtype=torch.float16, cache_enabled=True
        ):
            y_preds = model(inputs)
    y_preds = y_preds.detach().to("cpu").numpy().astype(np.float32)
    y_preds=  torch.softmax(torch.tensor(y_preds), 1).numpy()

    result = np.argmax(y_preds)

    if result == 0:
        return "0. llama2-70b-chat"
    elif result == 1:
        return "1. wizardLM-13b"
    elif result == 2:
        return "2. llama2-13b-chat"
    elif result == 3:
        return "3. wizardLM-70b"
    elif result == 4:
        return "4. llama2-7b-chat"
    elif result == 5:
        return "5. tinyllama-1b-chat"
    elif result == 6:
        return "6. mistral-7b-openorca"
    else:
        return "Error"




def do_submit(question, response):
    full_text = question + " " + response
    result = do_inference(full_text)
    return result

@spaces.GPU
def greet():
    pass

with gr.Blocks(title=title) as demo: # theme=theme
    sample_examples = pd.read_csv('sample_examples.csv')
    example_list = sample_examples[['Question','Response','target']].sample(2).values.tolist()
    gr.Markdown(f"## {title}")
    with gr.Row():
            question_text = gr.Textbox(lines=2, placeholder="Question:", label="")
            response_text = gr.Textbox(lines=2, placeholder="Response:", label="")
            target_text = gr.Textbox(lines=1, placeholder="Target:", label="", interactive=False , visible=False)
            llm_num = gr.Textbox(value="", label="LLM #")
    with gr.Row():
            sub_btn = gr.Button("Submit")
            sub_btn.click(fn=do_submit,  inputs=[question_text, response_text], outputs=[llm_num])

    gr.Markdown("## Sample Inputs:")
    gr.Examples(
        example_list,
        [question_text,response_text,target_text],
    )

demo.launch(greet)