Upload Orpheus_Auto_Continuations_Generator.ipynb

inference_code/Orpheus_Auto_Continuations_Generator.ipynb

@@ -344,6 +344,7 @@
     "temperature = 1.0\n",
     "top_p_value = 0.96\n",
     "num_mem_tokens = 7168 # up to 12 chunks\n",
+    "use_prime_embeddings = False\n",
     "\n",
     "# Advanced options\n",
     "max_tok_rep_ratio = 0.95\n",
@@ -377,9 +378,9 @@
     "    \n",
     "    inp = torch.LongTensor(y).cuda()\n",
     "    \n",
-    "    embs = get_embeddings(inp)\n",
+    "    prime_embs = get_embeddings(inp)\n",
     "    \n",
-    "    scores = cosine_similarity(embeddings, embs).max(axis=0)\n",
+    "    scores = cosine_similarity(embeddings, prime_embs).max(axis=0)\n",
     "\n",
     "    scores = [o for o in scores if o != max(scores)]\n",
     "\n",
@@ -484,9 +485,9 @@
     "        \n",
     "        inp = torch.LongTensor([song]).cuda()\n",
     "        \n",
-    "        embs = get_embeddings(inp)\n",
+    "        prime_embs = get_embeddings(inp)\n",
     "        \n",
-    "        start_score = cosine_similarity(embeddings, embs).max(axis=0)[0]\n",
+    "        start_score = cosine_similarity(embeddings, prime_embs).max(axis=0)[0]\n",
     "\n",
     "        b_size = batch_size\n",
     "        stop = False\n",
@@ -529,8 +530,13 @@
     "                    inp = torch.LongTensor(output).cuda()\n",
     "                    \n",
     "                    embs = get_embeddings(inp)\n",
-    "                    \n",
-    "                    scores = cosine_similarity(embeddings, embs).max(axis=0)\n",
+    "\n",
+    "                    if use_prime_embeddings:\n",
+    "                        scores = cosine_similarity(prime_embs, embs).max(axis=0)    \n",
+    "\n",
+    "                    else:\n",
+    "                        scores = cosine_similarity(embeddings, embs).max(axis=0)\n",
+    "                        \n",
     "                    output_scores.extend(scores)\n",
     "            \n",
     "                    scores = [o for o in output_scores if o != max(output_scores)]\n",