Ingest data from csv file and apply vector embeddings using custom python processor

 1import datetime
 2import hashlib
 3import json
 4
 5import ollama
 6import pytz
 7
 8from nifiapi.flowfiletransform import FlowFileTransform, FlowFileTransformResult
 9from nifiapi.properties import PropertyDescriptor, StandardValidators
10
11
12class OllamaEmbeddingsGenerator(FlowFileTransform):
13    class Java:
14        implements = ['org.apache.nifi.python.processor.FlowFileTransform']
15
16    class ProcessorDetails:
17        version = '0.0.1-SNAPSHOT'
18        dependencies = ['ollama', 'pytz']
19
20    def __init__(self, **kwargs):
21        super().__init__()
22        options = [
23            "llama3.2",
24            "nomic-embed-text",
25            "granite-embedding",
26            "snowflake-arctic-embed2",
27            "all-minilm"
28        ]
29        self.embeddings = PropertyDescriptor(name="Embeddings Name",
30                                             description="Name of the embeddings Model",
31                                             required=True,
32                                             allowable_values=options,
33                                             validators=[StandardValidators.NON_EMPTY_VALIDATOR])
34        self.descriptors = [self.embeddings]
35
36    def getPropertyDescriptors(self):
37        return self.descriptors
38
39    def transform(self, context, flowfile):
40        flow_file_content = flowfile.getContentsAsBytes().decode('utf-8', "ignore")
41        json_content = json.loads(flow_file_content)
42        title = json_content['title']
43        brand = json_content['brand']
44        embeddings_name = context.getProperty(self.embeddings.name).getValue()
45        self.logger.info(f"Generating Ollama Embeddings Using {embeddings_name} model")
46        self.logger.info(f"Embedding Text {title} and {brand}")
47        input_text = f"Title: {title}, Brand Name: {brand}"
48        json_content['text'] = input_text
49        json_content['content'] = input_text
50        json_content['metadata'] = {}
51        json_content['embedding'] = self.__get_ollama_embeddings(embeddings_name, input_text)
52        md5hash = hashlib.md5(json_content['asin'].encode('utf-8')).hexdigest()
53        json_content['id'] = md5hash
54        json_content['@@timestamp'] = datetime.datetime.now(pytz.timezone('UTC')).isoformat()
55
56        return FlowFileTransformResult(relationship="success", contents=json.dumps(json_content),
57                                       attributes=flowfile.getAttributes())
58
59    def __get_ollama_embeddings(self, embeddings_name: str, input_text: str):
60        ollama_response = ollama.embed(model=embeddings_name, input=input_text)
61        self.logger.info(f"Embeddings Text::::: length {len(ollama_response['embeddings'])}")
62        return ollama_response['embeddings'][0]

1json_content['embedding'] = self.__get_ollama_embeddings(embeddings_name, input_text)

1options = [
2        "llama3.2",
3        "nomic-embed-text",
4        "granite-embedding",
5        "snowflake-arctic-embed2",
6        "all-minilm"
7]