diff --git a/docs/src/notebooks/youtube_transcript_search.ipynb b/docs/src/notebooks/youtube_transcript_search.ipynb index e3cac08a..bd6f2bdc 100644 --- a/docs/src/notebooks/youtube_transcript_search.ipynb +++ b/docs/src/notebooks/youtube_transcript_search.ipynb @@ -36,7 +36,7 @@ } ], "source": [ - "!pip install --quiet openai datasets \n", + "!pip install --quiet openai datasets\n", "!pip install --quiet -U lancedb" ] }, @@ -213,7 +213,7 @@ "if \"OPENAI_API_KEY\" not in os.environ:\n", " # OR set the key here as a variable\n", " os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n", - " \n", + "\n", "client = OpenAI()\n", "assert len(client.models.list().data) > 0" ] @@ -234,9 +234,12 @@ "metadata": {}, "outputs": [], "source": [ - "def embed_func(c): \n", + "def embed_func(c):\n", " rs = client.embeddings.create(input=c, model=\"text-embedding-ada-002\")\n", - " return [rs.data[0].embedding]" + " return [\n", + " data.embedding\n", + " for data in rs.data\n", + " ]" ] }, { @@ -514,7 +517,7 @@ " prompt_start +\n", " \"\\n\\n---\\n\\n\".join(context.text) +\n", " prompt_end\n", - " ) \n", + " )\n", " return prompt" ] },