|
| 1 | +from pinecone import Pinecone |
| 2 | +from pinecone.grpc import PineconeGRPC |
| 3 | + |
| 4 | + |
| 5 | +class TestInferencePluginRerank: |
| 6 | + def test_rerank(self, api_key): |
| 7 | + pc = Pinecone(api_key=api_key) |
| 8 | + |
| 9 | + model = "bge-reranker-v2-m3" |
| 10 | + result = pc.inference.rerank( |
| 11 | + model=model, |
| 12 | + query="i love dogs", |
| 13 | + documents=["dogs are pretty cool", "everyone loves dogs", "I'm a cat person"], |
| 14 | + top_n=1, |
| 15 | + return_documents=True, |
| 16 | + ) |
| 17 | + assert len(result.data) == 1 |
| 18 | + assert result.data[0].index == 1 |
| 19 | + assert result.data[0].document.text == "everyone loves dogs" |
| 20 | + assert result.model == model |
| 21 | + assert isinstance(result.usage.rerank_units, int) |
| 22 | + assert result.usage.rerank_units == 1 |
| 23 | + |
| 24 | + def test_rerank_grpc(self, api_key): |
| 25 | + pc = PineconeGRPC(api_key=api_key) |
| 26 | + |
| 27 | + model = "bge-reranker-v2-m3" |
| 28 | + result = pc.inference.rerank( |
| 29 | + model=model, |
| 30 | + query="i love dogs", |
| 31 | + documents=["dogs are pretty cool", "everyone loves dogs", "I'm a cat person"], |
| 32 | + top_n=1, |
| 33 | + return_documents=True, |
| 34 | + ) |
| 35 | + assert len(result.data) == 1 |
| 36 | + assert result.data[0].index == 1 |
| 37 | + assert result.data[0].document.text == "everyone loves dogs" |
| 38 | + assert result.model == model |
| 39 | + assert isinstance(result.usage.rerank_units, int) |
| 40 | + assert result.usage.rerank_units == 1 |
0 commit comments