Updated model name.

Author: joaomonteirof

embedding_sandbox.ipynb

@@ -24,7 +24,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "OUTPUT_FOLDER = \"/mnt/colab_public/datasets/joao/mteb_results/\"\n",
+    "OUTPUT_FOLDER = \"./mteb_results/\"\n",
     "DEVICE = \"cuda:0\"\n",
     "BATCH_SIZE=32\n",
     "MAX_INPUT_LEN = 10000\n",
@@ -57,7 +57,7 @@
     "        super().__init__()\n",
     "\n",
     "        self.model_name = model_name\n",
-    "        self.tokenizer = prepare_tokenizer(model_name)\n",
+    "        self.tokenizer = prepare_tokenizer(model_name, use_auth_token=True)\n",
     "        self.encoder = AutoModel.from_pretrained(model_name, use_auth_token=True).to(DEVICE).eval()\n",
     "        self.device = device\n",
     "        self.max_input_len = max_input_len\n",
@@ -88,10 +88,10 @@
     "\n",
     "        return [emb.squeeze().numpy() for emb in input_sentences_embedding]\n",
     "\n",
-    "class BigCodeEncoder(BaseEncoder):\n",
+    "class StarEncoder(BaseEncoder):\n",
     "\n",
     "    def __init__(self, device, max_input_len, maximum_token_len):\n",
-    "        super().__init__(device, max_input_len, maximum_token_len, model_name = \"bigcode/bigcode-encoder\")\n",
+    "        super().__init__(device, max_input_len, maximum_token_len, model_name = \"bigcode/starencoder\")\n",
     "    \n",
     "    def forward(self, input_sentences):\n",
     "\n",
@@ -152,7 +152,7 @@
    ],
    "source": [
     "codebert = CodeBERT(DEVICE, MAX_INPUT_LEN, MAX_TOKEN_LEN)\n",
-    "bigcode_model = BigCodeEncoder(DEVICE, MAX_INPUT_LEN, MAX_TOKEN_LEN)"
+    "starencoder = StarEncoder(DEVICE, MAX_INPUT_LEN, MAX_TOKEN_LEN)"
    ]
   },
   {
@@ -167,7 +167,7 @@
     "]\n",
     "\n",
     "codebert_embeddings = codebert.encode(input_sentences)\n",
-    "bigcode_model_embeddings = bigcode_model.encode(input_sentences)\n"
+    "starencoder_embeddings = starencoder.encode(input_sentences)\n"
    ]
   },
   {
@@ -202,13 +202,13 @@
     }
    ],
    "source": [
-    "results_bigcode_encoder = evaluation.run(\n",
-    "    bigcode_model, \n",
-    "    output_folder=os.path.join(OUTPUT_FOLDER, \"bigcode_encoder\"), \n",
+    "results_starencoder = evaluation.run(\n",
+    "    starencoder, \n",
+    "    output_folder=os.path.join(OUTPUT_FOLDER, \"starencoder\"), \n",
     "    batch_size=BATCH_SIZE, \n",
     "    overwrite_results=True,)\n",
     "\n",
-    "results_bigcode_encoder"
+    "results_starencoder"
    ]
   },
   {