Repository: jerryjliu/gpt_index
File: llama-index-core/llama_index/core/base/embeddings/base.py (Python)

Code before:
                    )
                )
                event_id = self.callback_manager.on_event_start(
                    CBEventType.EMBEDDING,
                    payload={EventPayload.SERIALIZED: self.to_dict()},
                )
                callback_payloads.append((event_id, cur_batch))

                if not self.embeddings_cache:
                    embeddings_coroutines.append(self._aget_text_embeddings(cur_batch))
                elif self.embeddings_cache is not None:
                    embeddings = [None for i in range(len(cur_batch))]
                    non_cached_texts: List[Tuple[int, str]] = []
                    for i, txt in enumerate(cur_batch):
                        cached_emb = self.embeddings_cache.get(
                            key=txt, collection="embeddings"
                        )
                        if cached_emb is not None:
                            cached_key = next(iter(cached_emb.keys()))
                            embeddings[i] = cached_emb[cached_key]
                        else:
                            non_cached_texts.append((i, txt))

                    local_cur_batch = cur_batch

                    async def batch_coroutine():
                        nonlocal local_cur_batch, non_cached_texts, embeddings

                        if len(non_cached_texts) > 0:
                            batch_embeddings = await self._aget_text_embeddings(
                                [x[1] for x in non_cached_texts]
                            )
                            for j, text_embedding in enumerate(batch_embeddings):
                                orig_i = non_cached_texts[j][0]
                                embeddings[orig_i] = text_embedding
                                await self.embeddings_cache.aput(
                                    key=local_cur_batch[orig_i],
                                    val={str(uuid.uuid4()): text_embedding},
                                    collection="embeddings",
                                )
                        return embeddings

                    embeddings_coroutines.append(batch_coroutine())

                cur_batch = []

        # Execute all coroutines
        if num_workers and num_workers > 1:
            nested_embeddings = await run_jobs(
                embeddings_coroutines,
                show_progress=show_progress,

Review comment (performance):
We need a type annotation here, I think it should be something like:
```suggestion
                    async def batch_coroutine() -> List[List[float]]:
```

Code after:
                    elif self.embeddings_cache is not None:
                        embeddings = self._get_text_embeddings_cached(cur_batch)
                    result_embeddings.extend(embeddings)
                    event.on_end(
                        payload={
                            EventPayload.CHUNKS: cur_batch,
                            EventPayload.EMBEDDINGS: embeddings,
                        },
                    )
                dispatcher.event(
                    EmbeddingEndEvent(
                        chunks=cur_batch,
                        embeddings=embeddings,
                    )
                )
                cur_batch = []

        return result_embeddings

    @dispatcher.span
    async def aget_text_embedding_batch(
        self, texts: List[str], show_progress: bool = False
    ) -> List[Embedding]:
        """Asynchronously get a list of text embeddings, with batching."""
        num_workers = self.num_workers

        model_dict = self.to_dict()
        model_dict.pop("api_key", None)

        cur_batch: List[str] = []
        embeddings_coroutines: List[Coroutine] = []
        callback_payloads: List[Tuple[str, List[str]]] = []

        # for idx, text in queue_with_progress:
        for idx, text in enumerate(texts):
            cur_batch.append(text)
            if idx == len(texts) - 1 or len(cur_batch) == self.embed_batch_size:
                # flush
                dispatcher.event(
                    EmbeddingStartEvent(
                        model_dict=model_dict,
                    )
                )
                event_id = self.callback_manager.on_event_start(
                    CBEventType.EMBEDDING,
                    payload={EventPayload.SERIALIZED: self.to_dict()},
                )
                callback_payloads.append((event_id, cur_batch))

                if not self.embeddings_cache:
                    embeddings_coroutines.append(self._aget_text_embeddings(cur_batch))