Skip to content
Snippets Groups Projects
Select Git revision
  • f28b94a997cf18ae37ab7dd36e0a88f6d7fd4e0b
  • main default protected
2 results

pipeline.py

Blame
  • run_mpt.py 1.06 KiB
    #https://huggingface.co/bigscience/bloomz-7b1-mt
    from transformers import AutoModelForCausalLM, AutoTokenizer
    import transformers
    model = transformers.AutoModelForCausalLM.from_pretrained(
      'mosaicml/mpt-7b-instruct',
      trust_remote_code=True
    )
    
    def main(result_path: str, corpus_path: str, model: str = 'mpt-7b-instruct', template_id: str = '0'):
        checkpoint = 'mosaicml/' + model
    
        tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
        llm = AutoModelForCausalLM.from_pretrained(checkpoint, trust_remote_code=True).to('cuda')
    
        def generate(input_string):
            inputs = tokenizer(input_string, return_tensors="pt").input_ids.to("cuda")
            outputs = llm.generate(inputs, max_new_tokens=32)
    
            generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
            return generated[len(input_string):].split('#')[0]
    
        import deft
        results = deft.run_inference(generate, corpus_path, deft.template_from_id(template_id))
        deft.write_results(results, result_path)
    
    if __name__ == '__main__':
        import fire
        fire.Fire(main)