openai

/

circuit-sparsity

like 208

Follow

OpenAI 37.4k

import sys
import os
import torch

from transformers import PreTrainedTokenizerFast, AutoModelForCausalLM

current_dir = os.path.dirname(os.path.abspath(__file__))


parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)

local_model_path = current_dir

try:
    tokenizer = PreTrainedTokenizerFast.from_pretrained(
        local_model_path, 
        trust_remote_code=True
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        
except Exception as e:
    exit()

try:
    model = AutoModelForCausalLM.from_pretrained(
        local_model_path,
        trust_remote_code=True,
        device_map="auto",
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
    )
except Exception as e:
    exit()

print("-" * 30)
prompt = "The future of AI is"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.no_grad():
    outputs = model.generate(
        **inputs, 
        max_new_tokens=30, 
        do_sample=True, 
        temperature=0.8,
        pad_token_id=tokenizer.eos_token_id
    )

print(tokenizer.decode(outputs[0], skip_special_tokens=True))
print("-" * 30)

Encountered exception while importing circuit_sparsity: No module named 'circuit_sparsity.tokenizer_simple2k'