Upload 3 files

falcon-app.py

@@ -0,0 +1,94 @@
+import gradio as gr
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer
+)
+from peft import PeftModel
+
+model_name = "tiiuae/falcon-7b"
+model_id = "personachat-finetuned-3000-steps"
+template = open("template.txt", "r").read()
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name,
+    trust_remote_code = True
+)
+base_model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map = "auto",
+    load_in_8bit = True,
+    trust_remote_code = True,
+    low_cpu_mem_usage = True
+)
+tuned_model = PeftModel.from_pretrained(
+    base_model,
+    model_id
+)
+
+def parse_response(encoded_output, user_input):
+    decoded_output = tokenizer.batch_decode(encoded_output)[0]
+    decoded_output = decoded_output.replace(user_input, "")
+    decoded_output = decoded_output.split("<|endoftext|>",1)[0].strip()
+    return decoded_output
+
+def generate(personality, user_input, state = {"base_state":[], "tune_state":[]}):
+    try:
+        personality = "\n".join(personality.split("."))
+    except: pass
+    state["base_state"].append(user_input)
+    state["tune_state"].append(user_input)
+    base_prompt = template.format(
+        personality = personality,
+        history = "\n".join(state["base_state"])
+    )
+    tune_prompt = template.format(
+        personality = personality,
+        history = "\n".join(state["tune_state"])
+    )
+    print("****************************")
+    print(base_prompt)
+    print("****************************")
+    print(tune_prompt)
+    print("****************************")
+    base_input_ids = tokenizer(base_prompt, return_tensors="pt").to("cuda")
+    tune_input_ids = tokenizer(tune_prompt, return_tensors="pt").to("cuda")
+    kwargs = dict({
+        "top_k": 0,
+        "top_p": 0.9,
+        "do_sample": True,
+        "temperature": 0.5,
+        "max_new_tokens": 50,
+        "repetition_penalty": 1.1,
+        "num_return_sequences": 1
+    })
+    base_model_response = parse_response(
+        base_model.generate(
+            input_ids = base_input_ids["input_ids"],
+            **kwargs
+        ),
+        base_prompt
+    )
+    tune_model_response = parse_response(
+        tuned_model.generate(
+            input_ids = tune_input_ids["input_ids"],
+            **kwargs
+        ),
+        tune_prompt
+    )
+    state["base_state"].append(base_model_response)
+    state["tune_state"].append(tune_model_response)
+    return base_model_response, tune_model_response, state
+
+gr.Interface(
+    fn = generate,
+    inputs = [
+        gr.Textbox(label = "user personality", place_holder = "Enter your personality"),
+        gr.Textbox(label = "user chat", place_holder = "Enter your message"),
+        "state"
+    ],
+    outputs = [
+        gr.Textbox(label = "base model response"),
+        gr.Textbox(label = "fine tuned model response"),
+        "state"
+    ],
+    theme = "gradio/seafoam"
+).launch(share = True)

falcon-finetune-personachat.py

@@ -0,0 +1,100 @@
+import torch, einops
+from datasets import load_dataset
+from peft import LoraConfig
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    AutoTokenizer,
+    TrainingArguments
+)
+from peft.tuners.lora import LoraLayer
+
+from trl import SFTTrainer
+
+template = """### Personality:
+{personality}
+
+### History:
+{history}
+
+### Response:
+"""
+
+model_name = "tiiuae/falcon-7b"
+dataset_name = "bavard/personachat_truecased"
+
+def create_and_prepare_model():
+    compute_dtype = getattr(torch, "float16")
+
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=compute_dtype,
+        bnb_4bit_use_double_quant=True,
+    )
+
+    # device_map={"": 0}
+    device_map="auto"
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name, quantization_config=bnb_config, device_map=device_map, trust_remote_code=True
+    )
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map=device_map, trust_remote_code=True)
+
+    peft_config = LoraConfig(
+        lora_alpha=16,
+        lora_dropout=0.1,
+        r=64,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=[
+            "query_key_value"
+        ],
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    tokenizer.pad_token = tokenizer.eos_token
+
+    return model, peft_config, tokenizer
+
+
+training_arguments = TrainingArguments(
+    output_dir="./results",
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=4,
+    optim="paged_adamw_32bit",
+    save_steps=1000,
+    logging_steps=10,
+    learning_rate=2e-4,
+    fp16=True,
+    max_grad_norm=0.3,
+    max_steps=10000,
+    warmup_ratio=0.03,
+    group_by_length=False,
+    lr_scheduler_type="constant",
+)
+
+dataset = load_dataset(dataset_name, split="train")
+model, peft_config, tokenizer = create_and_prepare_model()
+model.config.use_cache = False
+
+def formatting_func(example):
+    return template.format(
+        personality = "\n".join(example["personality"]),
+        history = "\n".join(example["history"]),
+        response = example["candidates"][-1]
+    )
+
+trainer = SFTTrainer(
+    model=model,
+    train_dataset=dataset,
+    peft_config=peft_config,
+    max_seq_length=512,
+    tokenizer=tokenizer,
+    args=training_arguments,
+    packing=True,
+    formatting_func=formatting_func
+)
+
+trainer.train()

requirements.txt

@@ -0,0 +1,10 @@
+bitsandbytes
+git+https://github.com/huggingface/transformers.git 
+git+https://github.com/huggingface/peft.git
+git+https://github.com/huggingface/accelerate.git
+datasets
+trl
+einops
+scipy
+nvitop
+gradio