Set up the Dgraph schema for BEIR documents. Args: hnsw_config: HNSW index configuration
(self, hnsw_config: HNSWConfig)
| 35 | print(f"Connected to Dgraph at {config.grpc_endpoint} (256MB gRPC buffers)") |
| 36 | |
| 37 | def setup_schema(self, hnsw_config: HNSWConfig): |
| 38 | """ |
| 39 | Set up the Dgraph schema for BEIR documents. |
| 40 | |
| 41 | Args: |
| 42 | hnsw_config: HNSW index configuration |
| 43 | """ |
| 44 | print("Setting up Dgraph schema...") |
| 45 | |
| 46 | schema = f""" |
| 47 | doc_id: string @index(hash) . |
| 48 | title: string @index(term, fulltext) . |
| 49 | text: string @index(fulltext) . |
| 50 | embedding: float32vector @index({hnsw_config.to_index_string()}) . |
| 51 | |
| 52 | type Document {{ |
| 53 | doc_id |
| 54 | title |
| 55 | text |
| 56 | embedding |
| 57 | }} |
| 58 | """ |
| 59 | |
| 60 | op = pydgraph.Operation(schema=schema) |
| 61 | self.client.alter(op) |
| 62 | print("Schema setup complete") |
| 63 | |
| 64 | def reset_database(self): |
| 65 | """Drop all data and schema.""" |