Language: Python
Label: vulnerable
CWE: CWE-502 \u2014 Deserialization of Untrusted Data
The product deserializes untrusted data without sufficiently verifying that the resulting data will be valid.

Code:
    def test_legacy_index_retriever_retrieve(self):
        n_docs = 1
        retriever = self.get_dummy_legacy_index_retriever()
        hidden_states = np.array(
            [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)], dtype=np.float32
        )
        retrieved_doc_embeds, doc_ids, doc_dicts = retriever.retrieve(hidden_states, n_docs=n_docs)
        self.assertEqual(retrieved_doc_embeds.shape, (2, n_docs, self.retrieval_vector_size))
        self.assertEqual(len(doc_dicts), 2)
        self.assertEqual(sorted(doc_dicts[0]), ["text", "title"])
        self.assertEqual(len(doc_dicts[0]["text"]), n_docs)
        self.assertEqual(doc_dicts[0]["text"][0], "bar")  # max inner product is reached with second doc
        self.assertEqual(doc_dicts[1]["text"][0], "foo")  # max inner product is reached with first doc
        self.assertListEqual(doc_ids.tolist(), [[1], [0]])