Custom functions in flow

1
from instabase.udf_utils.clients.udf_helpers import get_output_ibmsg
2
3
def map_udf_func(input_record, step_folder, *args, **kwargs):
4
  input_filepath = input_record['input_filepath']
5
  file_content = input_record['content']
6
  output_filename = input_record['output_filename']
7
8
  # Copy output ibmsg from input
9
  output_ibmsg, err = get_output_ibmsg(input_filepath, step_folder,
10
                                       file_content)
11
12
  return {
13
    "out_files": [
14
      {
15
        "filename": f"map-udf-{output_filename}",
16
        "content": output_ibmsg
17
      }
18
    ]
19
  }
20
21
def register(name_to_fn):
22
  name_to_fn.update({
23
    'map_udf': {
24
      'fn': map_udf_func
25
    }
26
  })

1
fn_context = kwargs.get('_FN_CONTEXT_KEY')
2
input_record, err = fn_context.get_by_col_name('INPUT_RECORD')
3
file_content = input_record['content']

1
import json
2
import logging
3
import os
4
from typing import Any, Dict, Generator, List, Text
5
6
from instabase.ocr.client.libs.ibocr import ParsedIBOCRBuilder, RefinedPhrase
7
from instabase.udf_utils.clients.udf_helpers import get_output_ibmsg
8
9
_SUMMARY_FILENAME = 'summary.json'
10
11
def get_by_col_name(kwargs, col) -> Any:
12
  val, err = kwargs['_FN_CONTEXT_KEY'].get_by_col_name(col)
13
  if err:
14
    raise Exception(f'{col} is not available in UDF context')
15
  return val
16
17
def generate_summary(*args: Any, **kwargs: Any) -> Generator[Dict, None, None]:
18
  logging.info('Started running Reduce UDF function')
19
20
  job_id = get_by_col_name(kwargs, 'JOB_ID')
21
  input_records = get_by_col_name(kwargs, 'INPUT_RECORDS')
22
  root_output_folder = get_by_col_name(kwargs, 'ROOT_OUTPUT_FOLDER')
23
  step_folder = get_by_col_name(kwargs, 'STEP_FOLDER')
24
  clients = get_by_col_name(kwargs, 'CLIENTS')
25
  summary_dict = {}
26
27
  for payload in input_records:
28
    input_filepath = payload['input_filepath']
29
    output_filename = payload['output_filename']
30
    content = payload['content']
31
32
    # Loading ibmsg so we can get records from it.
33
    builder, err = ParsedIBOCRBuilder.load_from_str(input_filepath, content)
34
    if err:
35
      raise Exception(err)
36
37
    # Pulls out ibmsg records only.
38
    for ibocr_record in builder.get_ibocr_records():
39
      raw_input_filepath = ibocr_record.get_document_path()
40
      if raw_input_filepath not in summary_dict:
41
        summary_dict[raw_input_filepath] = []
42
      result = {}
43
      if ibocr_record.has_class_label():
44
        result['class_label'] = ibocr_record.get_class_label()
45
        if ibocr_record.has_classify_page_range():
46
          result['page_range'] = ibocr_record.get_classify_page_range()
47
        else:
48
          result['page_range'] = {}
49
          result['page_range']['start_page'] = ibocr_record.get_page_numbers()[0]+1
50
          result['page_range']['end_page'] = ibocr_record.get_page_numbers()[-1]+1
51
52
      result['extracted_fields'] = {}
53
      refined_phrases, _ = ibocr_record.get_refined_phrases()
54
      for phrase in refined_phrases:
55
        name = phrase.get_column_name()
56
        value = phrase.get_column_value()
57
        result['extracted_fields'][name] = value
58
      summary_dict[raw_input_filepath].append(result)
59
60
    # Copy output ibmsg from input.
61
    output_ibmsg, err = get_output_ibmsg(input_filepath, step_folder, content)
62
    if err:
63
      raise Exception(err)
64
65
    output_ibmsg_dict = {
66
        'out_files': [{
67
            'filename': output_filename,
68
            'content': output_ibmsg
69
        }]
70
    }
71
    yield output_ibmsg_dict
72
73
  out_path = os.path.join(root_output_folder, _SUMMARY_FILENAME)
74
  clients.ibfile.write_file(out_path, json.dumps(summary_dict))
75
76
  return
77
78
def register(name_to_fn):
79
  name_to_fn.update({
80
    'generate_summary': {
81
      'fn': generate_summary
82
    }
83
  })

1
from typing import Any, Text
2
import json
3
import time
4
5
def flow_info(clients: Any, root_output_folder:Text, job_id: Text, **kwargs):
6
  flow_info = {
7
      'Job ID': job_id,
8
      'Start timestamp': time.time(),
9
  }
10
  out_path = root_output_folder + '/flowinfo.json'
11
  clients.ibfile.write_file(out_path, json.dumps(flow_info))
12
  return
13
14
def register(name_to_fn):
15
  name_to_fn.update({
16
    'flow_info': {
17
      'fn': flow_info
18
    }
19
  })

1
import logging
2
import os
3
import json
4
5
from instabase.ocr.client.libs.ibocr import ParsedIBOCRBuilder
6
7
def write_summary(**kwargs):
8
  summary_dict = {}
9
10
  fn_context = kwargs.get('_FN_CONTEXT_KEY')
11
  clients, _ = fn_context.get_by_col_name('CLIENTS')
12
  root_out_folder, _ = fn_context.get_by_col_name('ROOT_OUTPUT_FOLDER')
13
14
  res_path = os.path.join(root_out_folder, 'batch.ibflowresults')
15
  output, err = clients.ibfile.read_file(res_path)
16
  if err:
17
    return None, err
18
19
  results = json.loads(output)
20
  if results['can_resume']:
21
    # can_resume = True implies flow is stopped at checkpoint.
22
    # If true, skip writing summary.
23
    return None, None
24
25
  for result_id in results['results']:
26
    result = results['results'][result_id]
27
    for record in result['records']:
28
      if record['status'] == 'OK':
29
        ibocr_path = record['ibocr_full_path']
30
        ibocr, err = clients.ibfile.read_file(ibocr_path)
31
        if err:
32
          return None, f'Failed to fetch ibocr path={ibocr_path} err={err}'
33
34
        builder, err = ParsedIBOCRBuilder.load_from_str(ibocr_path, ibocr)
35
        if err:
36
          return None, f'Failed to parse ibocr path={ibocr_path} err={err}'
37
38
        # Iterate over the IBOCR records.
39
        for ibocr_record in builder.get_ibocr_records():
40
          raw_input_filepath = ibocr_record.get_document_path()
41
          if raw_input_filepath not in summary_dict:
42
            summary_dict[raw_input_filepath] = []
43
          result = {}
44
          if ibocr_record.has_class_label():
45
            result['class_label'] = ibocr_record.get_class_label()
46
            if ibocr_record.has_classify_page_range():
47
              result['page_range'] = ibocr_record.get_classify_page_range()
48
            else:
49
              result['page_range'] = {}
50
              result['page_range']['start_page'] = ibocr_record.get_page_numbers()[0]+1
51
              result['page_range']['end_page'] = ibocr_record.get_page_numbers()[-1]+1
52
53
          result['extracted_fields'] = {}
54
          refined_phrases, _ = ibocr_record.get_refined_phrases()
55
          for phrase in refined_phrases:
56
            name = phrase.get_column_name()
57
            value = phrase.get_column_value()
58
            result['extracted_fields'][name] = value
59
          summary_dict[raw_input_filepath].append(result)
60
61
  _, err = clients.ibfile.write_file(os.path.join(root_out_folder, 'summary.json'), json.dumps(summary_dict, indent=4))
62
  if err:
63
    return None, f'Failed to write summary file err={err}'
64
  return summary_dict, None
65
66
67
def send_results(**kwargs):
68
  summary, err = write_summary(**kwargs)
69
  if err:
70
    logging.error(err)
71
    return
72
73
  if not summary:
74
    return
75
76
  # TODO: Send result to downstream system
77
  pass
78
79
def register(name_to_fn):
80
  name_to_fn.update({
81
    'send_results': {
82
      'fn': send_results
83
    }
84
  })