This repository provides a GPT-NeoX based model with 1.4B parameters pre-trained on Japanese corpus of about 20B tokens. This model is developed by Stockmark Inc.
How to use
importtorchfromtransformersimportAutoModelForCausalLM,AutoTokenizer# Use torch.bfloat16 for A100 GPU and torch.flaot16 for the older generation GPUstorch_dtype=torch.bfloat16iftorch.cuda.is_available()andhasattr(torch.cuda,"is_bf16_supported")andtorch.cuda.is_bf16_supported()elsetorch.float16model=AutoModelForCausalLM.from_pretrained("stockmark/gpt-neox-japanese-1.4b",device_map="auto",torch_dtype=torch_dtype)tokenizer=AutoTokenizer.from_pretrained("stockmark/gpt-neox-japanese-1.4b")inputs=tokenizer("自然言語処理は",return_tensors="pt").to(model.device)withtorch.no_grad():tokens=model.generate(**inputs,max_new_tokens=128,repetition_penalty=1.1)output=tokenizer.decode(tokens[0],skip_special_tokens=True)print(output)