Question/Answer with xgen

[tonyphoang]

Hi all,

I've been able to get xgen7b to work with sentence completion using GPU but cannot get it to work with question/answer.

The code I'm using is below:

import torch
torch.cuda.empty_cache()
from transformers import AutoTokenizer, AutoModelForCausalLM
device_map = {"cuda:0" if torch.cuda.is_available() else "cpu"}
model_name='Salesforce/xgen-7b-8k-base'

tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model =AutoModelForCausalLM.from_pretrained(model_name, 
                                            torch_dtype=torch.bfloat16, 
                                            device_map="auto")

header = (
    "You are an artificial intelligence assistant. "
    "The assistant gives helpful, detailed, and polite answers to the human's questions.\n\n"
)
prompt = f"who is the president of the usa?"

inputs = tokenizer(header + prompt, return_tensors="pt").to('cuda')
sample = model.generate(**inputs, max_length=128, do_sample=True, top_k=100, eos_token_id=50256)
print(tokenizer.decode(sample[0]))