|
| 1 | +```python |
| 2 | +import os |
| 3 | + |
| 4 | +from unstructured.ingest.connector.astra import ( |
| 5 | + AstraAccessConfig, |
| 6 | + AstraWriteConfig, |
| 7 | + SimpleAstraConfig, |
| 8 | +) |
| 9 | +from unstructured.ingest.connector.local import SimpleLocalConfig |
| 10 | +from unstructured.ingest.interfaces import ( |
| 11 | + ChunkingConfig, |
| 12 | + EmbeddingConfig, |
| 13 | + PartitionConfig, |
| 14 | + ProcessorConfig, |
| 15 | + ReadConfig, |
| 16 | +) |
| 17 | +from unstructured.ingest.runner import LocalRunner |
| 18 | +from unstructured.ingest.runner.writers.astra import ( |
| 19 | + AstraWriter, |
| 20 | +) |
| 21 | +from unstructured.ingest.runner.writers.base_writer import Writer |
| 22 | + |
| 23 | + |
| 24 | +def get_writer() -> Writer: |
| 25 | + return AstraWriter( |
| 26 | + connector_config=SimpleAstraConfig( |
| 27 | + access_config=AstraAccessConfig( |
| 28 | + token=os.getenv("ASTRA_DB_TOKEN"), api_endpoint=os.getenv("ASTRA_DB_ENDPOINT") |
| 29 | + ), |
| 30 | + collection_name="test_collection", |
| 31 | + embedding_dimension=384, |
| 32 | + ), |
| 33 | + write_config=AstraWriteConfig(batch_size=80), |
| 34 | + ) |
| 35 | + |
| 36 | + |
| 37 | +if __name__ == "__main__": |
| 38 | + writer = get_writer() |
| 39 | + runner = LocalRunner( |
| 40 | + processor_config=ProcessorConfig( |
| 41 | + verbose=True, |
| 42 | + output_dir="local-output-to-astra", |
| 43 | + num_processes=2, |
| 44 | + ), |
| 45 | + connector_config=SimpleLocalConfig( |
| 46 | + input_path="example-docs/book-war-and-peace-1p.txt", |
| 47 | + ), |
| 48 | + read_config=ReadConfig(), |
| 49 | + partition_config=PartitionConfig(), |
| 50 | + chunking_config=ChunkingConfig(chunk_elements=True), |
| 51 | + embedding_config=EmbeddingConfig( |
| 52 | + provider="langchain-huggingface", |
| 53 | + ), |
| 54 | + writer=writer, |
| 55 | + writer_kwargs={}, |
| 56 | + ) |
| 57 | + runner.run() |
| 58 | +``` |
0 commit comments