ntdas/

pygmalion_6B_yaya

        from transformers import  GPTJForCausalLM,AutoTokenizer, GenerationConfig
        from peft import PeftModel

        lora_weights = 'kietbs/pygmalion_6B_yaya'  # Please download the weight, and change this path accordingly
        load_in_8bit = True
        model = GPTJForCausalLM.from_pretrained(pretrain_name, load_in_8bit=load_in_8bit, device_map='auto', torch_dtype=torch.float16)
        model = PeftModel.from_pretrained(model,lora_weights,torch_dtype=torch.float16,device_map={'':0})
        model = torch.compile(model)
        tokenizer = AutoTokenizer.from_pretrained('pygmalion-6b') #The orginal pretrained
        gen_config=GenerationConfig(
                temperature=0.1,
                top_p=0.75,
                top_k=40,
                num_beams=4
            )

        text = '[User]: What's the best food in Hanoi?''
        input_ids = tokenizer(text, return_tensors='pt')['input_ids'].to('cuda')
        with torch.no_grad():
            output = model.generate(input_ids=input_ids, generation_config=gen_config,return_dict_in_generate=True, output_scores=True,max_new_tokens=256)
            s = output.sequences[0]
            output = tokenizer.decode(s)
            print('Raw:',output)