Code for training and inference

README.md

@@ -12,6 +12,162 @@ tags:
 
 # Bloom 560M Finetuned on Instructions
 
+## Credit
+
+Code 99.99% copied from 
+*https://github.com/bofenghuang/vigogne*
+*https://colab.research.google.com/drive/1jCkpikz0J2o20FBQmYmAGdiKmJGOMo-o?usp=sharing#scrollTo=DpYr24pR8T_0*
+
+
+# Inference Code
+
+```python
+
+from peft import PeftModel
+from transformers import PreTrainedTokenizer, PreTrainedModel, AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModelForCausalLM, LoraConfig
+from typing import Optional
+from transformers import GenerationConfig
+import torch
+
+PROMPT_DICT = {
+    "prompt_input": (
+        "Below is a^n instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n"
+        "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response:\n"
+    ),
+    "prompt_no_input": (
+        "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n"
+        "### Instruction:\n{instruction}\n\n### Response:\n"
+    ),
+}
+
+
+def get_model(model_name_or_path: str, load_in_8bit: bool = True, device_map="auto",
+              cpu: bool = False) -> PreTrainedModel:
+    if cpu:
+        model = AutoModelForCausalLM.from_pretrained(model_name_or_path, device_map=device_map,
+                                                     low_cpu_mem_usage=True)
+    else:
+        model = AutoModelForCausalLM.from_pretrained(model_name_or_path, load_in_8bit=load_in_8bit,
+                                                     device_map=device_map, torch_dtype=torch.float16)
+
+    return model
+
+
+def get_peft_model(model: PreTrainedModel, lora_model_name_or_path: Optional[str] = None) -> PeftModelForCausalLM:
+    model = PeftModel.from_pretrained(model, lora_model_name_or_path, torch_dtype=torch.float16)
+
+    return model
+
+
+def get_tokenizer(model_name_or_path: str, max_input_len: int) -> PreTrainedTokenizer:
+    tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, model_max_length=max_input_len,
+                                              padding_side="right", use_fast=False)
+
+    return tokenizer
+
+
+def get_llm_inference_model(base_model_name_or_path: str, lora_model_name_or_path: str, load_in_8bit: bool,
+                            device_map) -> PeftModel:
+    cpu = True if not torch.cuda.is_available() else False
+
+    model = get_model(base_model_name_or_path, load_in_8bit, device_map, cpu=cpu)
+
+    model = get_peft_model(model, lora_model_name_or_path=lora_model_name_or_path)
+
+    if not load_in_8bit:
+        model.half()
+
+    model.eval()
+
+    if torch.__version__ >= "2":
+        model = torch.compile(model)
+
+    return model
+
+
+def generate_prompt(example):
+    return (
+        PROMPT_DICT["prompt_input"].format_map(example)
+        if example["input"]
+        else PROMPT_DICT["prompt_no_input"].format_map(example)
+    )
+
+
+def infer(instruction: str, input_text: Optional[str] = None, temperature: float = 0.1, top_p: float = 0.95,
+          max_new_tokens: int = 512, early_stopping: bool = True, do_sample: bool = True,
+          repetition_penalty: float = 2.5) -> str:
+    prompt = generate_prompt({"instruction": instruction, "input": input_text})
+
+    tokenized_inputs = tokenizer(prompt, return_tensors="pt")
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    input_ids = tokenized_inputs["input_ids"].to(device)
+
+    generation_config = GenerationConfig(temperature=temperature, top_p=top_p, do_sample=do_sample,
+                                         repetition_penalty=repetition_penalty, early_stopping=early_stopping)
+
+    with torch.inference_mode():
+        generation_output = model.generate(input_ids=input_ids, generation_config=generation_config,
+                                           return_dict_in_generate=True, max_new_tokens=max_new_tokens)
+
+    output = generation_output.sequences[0]
+
+    output = tokenizer.decode(output, skip_special_tokens=True)
+
+    return output.split("### Response:")[1].strip()
+
+
+base_model_name_or_path = "bigscience/bloom-560m"
+
+lora_model_name_or_path = "crayon-coe/alpaca-bloom-560m-en"
+
+model = get_llm_inference_model(base_model_name_or_path, lora_model_name_or_path, True, "auto")
+
+tokenizer = get_tokenizer(base_model_name_or_path, 512)
+
+context = "Write a letter expressing your love for computers"
+
+output = infer(context)
+
+print(output)
+
+# Output
+# I am so grateful to have been able access this wonderful computer system and its amazing features, which I can now use daily with ease. 
+# 
+# My heartfelt thanks go out in advance of all my friends who are using it as well. 
+# Thank you again!
+
+```
+
+# Training Parameters
+
+```json
+{
+    "max_input_len": 512,
+    "load_in_8bit": True,
+    "model_name_or_path": "bigscience/bloom-560m",
+    "device_map": "auto",
+    "bias": "none",
+    "lora_dropout": 0.05,
+    "lora_alpha": 32,
+    "target_modules": ["query_key_value"],
+    "task_type": "CAUSAL_LM",
+    "lora_r": 16,
+    "pad_to_multiple_of": 8,
+    "num_train_epochs": 3,
+    "learning_rate": 0.0003,
+    "gradient_accumulation_steps": 16,
+    "per_device_train_batch_size": 8,
+    "val_set_size": 500,
+    "save_steps": 200,
+    "eval_steps": 200,
+    "evaluation_strategy": "steps",
+    "save_strategy": "steps"
+}
+```
+
 # Training Code
 
 ```python