sarvamai

/

OpenHathi-7B-Hi-v0.1-Base

like 105

Follow

Sarvam AI 140

If you are struggling to run it on Google Colab T4 or a single T4 GPU system with low CPU in general, you can achieve through the following steps:

1. Along with torch and transformers make sure to install accelerate

pip install accelerate

2. Use cuda device configuration to run it on GPU

import torch
from transformers import LlamaTokenizer, LlamaForCausalLM

# setting GPU as default device 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# loading model and tokenizer
hf_model = 'sarvamai/OpenHathi-7B-Hi-v0.1-Base'
tokenizer = LlamaTokenizer.from_pretrained(hf_model)
model = LlamaForCausalLM.from_pretrained(hf_model, torch_dtype=torch.bfloat16, device_map=device)

# sample inference
prompt = "मैं एक अच्छा हाथी हूँ"
input_tokens = tokenizer(prompt, return_tensors="pt").to(device)
output_tokens = model.generate(input_tokens.input_ids, max_length=256)[0]
output = tokenizer.decode(output_tokens, skip_special_tokens=True)
print(output)

Here is the Google Colab link if you want to get started quickly.