## Base imports
from dataiku.code_env_resources import clear_all_env_vars
from dataiku.code_env_resources import set_env_path
from dataiku.code_env_resources import set_env_var
from dataiku.code_env_resources import grant_permissions
from transformers import AutoTokenizer, AutoModel

# Clears all environment variables defined by previously run script
clear_all_env_vars()

## Hugging Face
# Set HuggingFace cache directory
set_env_path("TRANSFORMERS_CACHE", "huggingface")

tokenizer = AutoTokenizer.from_pretrained("DeepChem/ChemBERTa-77M-MLM", cache_dir="huggingface")
model = AutoModel.from_pretrained("DeepChem/ChemBERTa-77M-MLM", cache_dir="huggingface")

grant_permissions("huggingface")