Processor model prototype

trustgraph/embeddings/vectorize/vectorize.py

@@ -16,18 +16,24 @@ from ... schema import Chunk, VectorsChunk
 from ... embeddings_client import EmbeddingsClient
 from ... log_level import LogLevel
 
+default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
+default_input_queue = 'chunk-load'
+default_output_queue = 'vectors-chunk-load'
+default_subscriber = 'embeddings-vectorizer'
+
 class Processor:
 
     def __init__(
             self,
-            pulsar_host,
-            input_queue,
-            output_queue,
-            subscriber,
-            log_level,
-            model,
+            pulsar_host=default_pulsar_host,
+            input_queue=default_input_queue,
+            output_queue=default_output_queue,
+            subscriber=default_subscriber,
+            log_level=LogLevel.INFO,
     ):
 
+        self.client = None
+        
         self.client = pulsar.Client(
             pulsar_host,
             logger=pulsar.ConsoleLogger(log_level.to_pulsar())
@@ -89,7 +95,9 @@ class Processor:
                 self.consumer.negative_acknowledge(msg)
 
     def __del__(self):
-        self.client.close()
+
+        if self.client:
+            self.client.close()
 
 def run():
 
@@ -98,11 +106,6 @@ def run():
         description=__doc__,
     )
 
-    default_pulsar_host = os.getenv("PULSAR_HOST", 'pulsar://pulsar:6650')
-    default_input_queue = 'chunk-load'
-    default_output_queue = 'vectors-chunk-load'
-    default_subscriber = 'embeddings-vectorizer'
-
     parser.add_argument(
         '-p', '--pulsar-host',
         default=default_pulsar_host,
@@ -135,12 +138,6 @@ def run():
         help=f'Output queue (default: info)'
     )
 
-    parser.add_argument(
-        '-m', '--model',
-        default="all-MiniLM-L6-v2",
-        help=f'LLM model (default: all-MiniLM-L6-v2)'
-    )
-
     args = parser.parse_args()
 
     while True:
@@ -153,7 +150,6 @@ def run():
                 output_queue=args.output_queue,
                 subscriber=args.subscriber,
                 log_level=args.log_level,
-                model=args.model,
             )
 
             p.run()