pubmedGPT動かしてみたメモ

これで動いた．
colabで動かそうと思ったけどモデルが重すぎてcolabでは動かなかったが，他の環境で動くことを確認した．
一応colab用のコードを貼っておく．

!wget -c https://repo.continuum.io/miniconda/Miniconda3-latest-Linux-x86_64.sh
!chmod +x Miniconda3-latest-Linux-x86_64.sh
!bash ./Miniconda3-latest-Linux-x86_64.sh -b -f -p /usr/local
!conda install -q -y -c conda-forge rdkit python=3.8
import sys
sys.path.append('/usr/local/lib/python3.8/site-packages')

!conda create -n pubmedgpt python=3.8.12 pytorch=1.12.1 torchdata cudatoolkit=11.3 -c pytorch
!conda activate pubmedgpt
!pip install -r finetune/setup/requirements.txt

!git clone https://github.com/stanford-crfm/pubmedgpt
%cd pubmedgpt/finetune # sys.path.appendの方が良いかも

あとはgithubにあるものと同じようにして動かせばOK

import torch

from transformers import GPT2LMHeadModel, GPT2Tokenizer

device = torch.device("cuda")

tokenizer = GPT2Tokenizer.from_pretrained("stanford-crfm/pubmed_gpt_tokenizer")

model = GPT2LMHeadModel.from_pretrained("stanford-crfm/pubmedgpt").to(device)

input_ids = tokenizer.encode(
    "Photosynthesis is ", return_tensors="pt"
).to(device)

sample_output = model.generate(input_ids, do_sample=True, max_length=50, top_k=50)

print("Output:\n" + 100 * "-")
print(tokenizer.decode(sample_output[0], skip_special_tokens=True))