diff --git a/embedding_sandbox.ipynb b/embedding_sandbox.ipynb index 6589d00..c07152e 100644 --- a/embedding_sandbox.ipynb +++ b/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" ] }, {