|
import torch |
|
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig |
|
|
|
from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training, get_peft_model |
|
|
|
model_id = "/share/models/open-zharfa" |
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_id, add_eos_token=True) |
|
|
|
base_model = AutoModelForCausalLM.from_pretrained( |
|
model_id, |
|
low_cpu_mem_usage=True, |
|
return_dict=True, |
|
torch_dtype=torch.float16, |
|
device_map="auto", |
|
) |
|
|
|
base_model.generation_config.do_sample = True |
|
|
|
|
|
|
|
tokenizer.pad_token = tokenizer.unk_token |
|
tokenizer.padding_side = "right" |
|
|
|
def get_completion_merged(query: str, model, tokenizer) -> str: |
|
device = "cuda:0" |
|
|
|
prompt_template = """ |
|
|
|
GPT4 Correct User: {query}<|end_of_turn|>GPT4 Correct Assistant: |
|
|
|
""" |
|
prompt = prompt_template.format(query=query) |
|
|
|
encodeds = tokenizer(prompt, return_tensors="pt", add_special_tokens=True) |
|
|
|
model_inputs = encodeds.to(device) |
|
|
|
generated_ids = model.generate(**model_inputs, max_new_tokens=1000, do_sample=True, temperature=0.5, pad_token_id=tokenizer.unk_token_id) |
|
decoded = tokenizer.batch_decode(generated_ids) |
|
return (decoded[0]) |
|
|
|
while True: |
|
q = input('q : ') |
|
result = get_completion_merged(query=q, model=base_model, tokenizer=tokenizer) |
|
print(result) |
|
|
