transformers_quick_loader.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup and activate virtual environment first\n",
    "\n",
    "pip install transformers    #add a \"!\" at the beginning if required"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModel, AutoTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prep_loader(model_name, size=None):\n",
    "    if size:\n",
    "        _model = AutoModel.from_pretrained(str(model_name) + str(size))\n",
    "        _tokenizer = AutoTokenizer.from_pretrained(str(model_name) + str(size))\n",
    "        \n",
    "    else:\n",
    "        _model = AutoModel.from_pretrained(str(model_name))\n",
    "        _tokenizer = AutoTokenizer.from_pretrained(str(model_name))\n",
    "\n",
    "    return _model, _tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_model_request = input(\"Please input required Model name, (optional)model size: \")\n",
    "\n",
    "model, tokenizer = prep_loader(user_model_request)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#To check if model has loaded properly\n",
    "print(model.num_parameters) "
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}