lyogavin/Anima-7B-100K

# Please update the path of `CUDA_HOME`
export CUDA_HOME=/usr/local/cuda-11.8
pip install transformers==4.31.0
pip install sentencepiece
pip install ninja
pip install flash-attn --no-build-isolation
pip install git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/rotary
pip install git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/xentropy
pip install evaluate
pip install git+https://github.com/huggingface/peft.git@v0.4.0
pip install wandb

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

base_model = "lyogavin/Anima-7B-100K"
tokenizer = AutoTokenizer.from_pretrained(base_model)
model = AutoModelForCausalLM.from_pretrained(
            base_model,
            torch_dtype=torch.float16,
            trust_remote_code=True,
            device_map="auto", 
        )
model.eval()

prompt = "Where is the capital of US?"
inputs = tokenizer(prompt, return_tensors="pt")

inputs['input_ids'] = inputs['input_ids'].cuda()
inputs['attention_mask'] = inputs['attention_mask'].cuda()

# Generate
generate_ids = model.generate(**inputs, max_new_tokens=30,
                       only_last_logit=True, # to save memory
                       use_cache=False, # when run into OOM, enable this can save memory
                       xentropy=True)
output = tokenizer.batch_decode(generate_ids, 
                                skip_special_tokens=True,
                                clean_up_tokenization_spaces=False)[0]

./run_longer_training.sh