phi2-prompt-recovery

LLM Prompt Recovery

This model contains a LoRA adapter on Phi-2 for predicting the LLM prompt that is used to convert text A to text B, created for Kaggle's LLM prompt recovery contest.

It was trained with rewritten text generation by Gemma 7B Quantized model.

Usage

Load model:

import bitsandbytes
import accelerate
import transformers
import optimum
import torch

import sys 
import contextlib
import os

from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoConfig
from functools import cache

from peft import PeftModel

device = torch.device("cuda")


PEFT_PATH = "osy-it/phi2-prompt-recovery"
MODEL_PATH = "microsoft/phi-2"
def load_model():
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, model_max_length=512)
    tokenizer.padding_side = "left"
    tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_PATH,
        device_map = "auto",
        trust_remote_code = True,
        torch_dtype=torch.float16,
    )
    
    model = PeftModel.from_pretrained(model, PEFT_PATH)
    model = model.merge_and_unload()
    tokenizer.padding_side = "left"
   
    return model, tokenizer

Predict prompt using the following chat template:

RECOVER_PROMPT = """Given these two texts above, infer the instruction that was given to convert the original text to the rewritten text. Make your best guess. Only output your inferred instruction."""
INSTRUCTION = """#Original Text\n```{original_text}```\n\n#Rewritten Text\n```{rewritten_text}```\n\nInstruct: {RECOVER_PROMPT}"""
CHAT = "{instruction}\nPrompt:"
response_template = "\nPrompt:"

def truncate(text, _len):
    words = text.split()
    if len(words) > _len:
        text = " ".join(words[:_len]+["[...]"])
    return text

def generate(samples, model):
  inputs = []
  for s in samples:
      original_text = truncate(s["original_text"], 100)
      rewritten_text = truncate(s["rewritten_text"], 200)
    
      i = INSTRUCTION.format(RECOVER_PROMPT=RECOVER_PROMPT, original_text=original_text, rewritten_text=rewritten_text)
      t = CHAT.format(instruction=i)
      inputs.append(t)
      
  
  tokens = tokenizer(inputs, return_tensors="pt", add_special_tokens=False, truncation=True, padding=True, max_length=512).to(device)

  in_len = tokens["input_ids"].shape[1]
    
  res = model.generate(max_length=in_len+50, eos_token_id=[tokenizer.eos_token_id], pad_token_id=tokenizer.eos_token_id, do_sample=False, num_beams=3, **tokens)

  # Post process to grab the first sentence
  results = tokenizer.batch_decode(res, skip_special_tokens=False)
  results = [r.split(response_template)[1] if response_template in r else "" for r in results ]
  results = [r.split("\n")[0] for r in results]

  del tokens,in_len
  return results