{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/haoyang/miniconda3/envs/llm_reason/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "from src.submission.check_validity import is_model_on_hub\n",
    "from huggingface_hub import HfApi\n",
    "import re"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.bfloat16"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "still_on_hub, _, model_config = is_model_on_hub(\n",
    "    \"01-ai/Yi-34B-Chat\", \"main\", trust_remote_code=True, test_tokenizer=False\n",
    ")\n",
    "getattr(model_config, \"torch_dtype\", None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Yi-34b is on hub: True\n",
      "Yi-34b config: torch.bfloat16\n",
      "Yi-34b size: 34.389\n",
      "Mistral-7b is on hub: True\n",
      "Mistral-7b config: torch.bfloat16\n",
      "Mistral-7b size: 7.242\n",
      "Vicuna-13b is on hub: True\n",
      "Vicuna-13b config: torch.float16\n",
      "Vicuna-13b size: 13.0\n",
      "Phi-1.5 is on hub: True\n",
      "Phi-1.5 config: torch.float16\n",
      "Phi-1.5 size: N/A\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/haoyang/.cache/huggingface/modules/transformers_modules/mosaicml/mpt-30b-instruct/56bcbea5361d8381c297ca51c02ee5b6f0415cb4/configuration_mpt.py:97: UserWarning: alibi is turned on, setting `learned_pos_emb` to `False.`\n",
      "  warnings.warn(f'alibi is turned on, setting `learned_pos_emb` to `False.`')\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MPT-30b is on hub: True\n",
      "MPT-30b config: torch.bfloat16\n",
      "MPT-30b size: 30.0\n",
      "Phi-2 is on hub: True\n",
      "Phi-2 config: torch.float16\n",
      "Phi-2 size: 2.78\n",
      "Qwen-14b is on hub: True\n",
      "Qwen-14b config: None\n",
      "Qwen-14b size: 14.167\n"
     ]
    }
   ],
   "source": [
    "open_models = {\n",
    "    \"Yi-34b\": \"01-ai/Yi-34B-Chat\",\n",
    "    \"Mistral-7b\": \"mistralai/Mistral-7B-Instruct-v0.1\",\n",
    "    \"Vicuna-13b\": \"lmsys/vicuna-13b-v1.3\",\n",
    "    \"Phi-1.5\": \"microsoft/phi-1_5\",\n",
    "    \"MPT-30b\": \"mosaicml/mpt-30b-instruct\",\n",
    "    \"Phi-2\": \"microsoft/phi-2\",\n",
    "    \"Qwen-14b\": \"Qwen/Qwen-14B-Chat\"\n",
    "}\n",
    "\n",
    "api = HfApi()\n",
    "size_pattern = size_pattern = re.compile(r\"(\\d\\.)?\\d+(b|m)\")\n",
    "\n",
    "new_params = {}\n",
    "\n",
    "for model_name, model_id in open_models.items():\n",
    "    still_on_hub, _, model_config = is_model_on_hub(\n",
    "        model_id, \"main\", trust_remote_code=True, test_tokenizer=False\n",
    "    )\n",
    "    precision = str(getattr(model_config, \"torch_dtype\", None))\n",
    "\n",
    "    print(f\"{model_name} is on hub: {still_on_hub}\")\n",
    "    print(f\"{model_name} config: {precision}\")\n",
    "\n",
    "    model_info = api.model_info(repo_id=model_id, revision=\"main\")\n",
    "    try:\n",
    "        model_size = round(model_info.safetensors[\"total\"] / 1e9, 3)\n",
    "    except (AttributeError, TypeError):\n",
    "        try:\n",
    "            size_match = re.search(size_pattern, model_info.modelId.lower())\n",
    "            model_size = size_match.group(0)\n",
    "            model_size = round(float(model_size[:-1]) if model_size[-1] == \"b\" else float(model_size[:-1]) / 1e3, 3)\n",
    "        except AttributeError:\n",
    "            model_size = \"N/A\"\n",
    "    print(f\"{model_name} size: {model_size}\")\n",
    "    new_params[model_name] = (precision, model_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'Yi-34b': ('torch.bfloat16', 34.389),\n",
       " 'Mistral-7b': ('torch.bfloat16', 7.242),\n",
       " 'Vicuna-13b': ('torch.float16', 13.0),\n",
       " 'Phi-1.5': ('torch.float16', 'N/A'),\n",
       " 'MPT-30b': ('torch.bfloat16', 30.0),\n",
       " 'Phi-2': ('torch.float16', 2.78),\n",
       " 'Qwen-14b': ('None', 14.167)}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "new_params"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm_reason",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}