Stopping criteria for text-to-sql task with text-generation pipeline

stopping_criteria.py

from transformers import AutoTokenizer, AutoModelForCausalLM, \
    GenerationConfig, BitsAndBytesConfig, StoppingCriteria, \
    TextStreamer, pipeline
import torch


class GenerateSqlStoppingCriteria(StoppingCriteria):

    def __call__(self, input_ids, scores, **kwargs):
        # stops when sequence "```\n" is generated
        # Baichuan2 tokenizer
        # ``` -> 84
        # \n -> 5
        return (
            len(input_ids[0]) > 1
            and input_ids[0][-1] == 5
            and input_ids[0][-2] == 84
        )
    def __len__(self):
        return 1
    
    def __iter__(self):
        yield self

        
model_id = "baichuan-inc/Baichuan2-13B-chat"
tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    use_fast=False,
    trust_remote_code=True,
    revision="v2.0"
)
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    quantization_config=quantization_config,
    trust_remote_code=True,
)
model.generation_config = GenerationConfig.from_pretrained(model_id, revision="v2.0")
streamer = TextStreamer(tokenizer, skip_prompt=True,)
pipeline = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    revision="v2.0",
    do_sample=False,
    num_return_sequences=1,
    eos_token_id=tokenizer.eos_token_id,
    stopping_criteria=GenerateSqlStoppingCriteria(),
    streamer=streamer,
)