AWSTemplateFormatVersion: "2010-09-09"
Transform: AWS::Serverless-2016-10-31

Parameters:
  InitialInstanceCount:
    Type: Number
    Description: Number of instances to launch initially. Specify 0 to provision a serverless endpoint.
    Default: 1
  InstanceType:
    Type: String
    Description: The ML compute instance type.
    Default: ml.m5.xlarge
  CallEventsTable:
    Type: String
    Description: Name of the LCA Call Events Table to look up transcripts.
  FetchTranscriptArn:
    Type: String
    Description: Arn of the Lambda function that fetches the transcript as a string.

Conditions:
  ServerlessSageMakerEndpoint: !Equals
    - !Ref 'InitialInstanceCount'
    - 0
  ProvisionedSageMakerEndpoint: !Not
    - !Equals
      - !Ref 'InitialInstanceCount'
      - 0

Outputs:
  EndpointName: 
    Description: Summarizer Model Endpoint Name
    Value: !If
      - ProvisionedSageMakerEndpoint
      - !GetAtt LCASMProvisionedSummarizationEndpoint.EndpointName
      - !GetAtt LCASMServerlessSummarizationEndpoint.EndpointName 
  InvokeLambdaArn: 
    Description: ARN of the Lambda that invokes the summarizer endpoint
    Value: !GetAtt SummaryLambda.Arn

Resources:
  # 763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference:1.7.1-transformers4.6.1-gpu-py36-cu110-ubuntu18.04
  LCASMSummarizationModel:
    Type: AWS::SageMaker::Model
    Properties:
      PrimaryContainer:
        # image resource found at https://github.com/aws/deep-learning-containers/blob/master/available_images.md
        # old: Image: !Sub 763104351884.dkr.ecr.${AWS::Region}.amazonaws.com/huggingface-pytorch-inference:1.10.2-transformers4.17.0-cpu-py38-ubuntu20.04
        Image: !Sub 763104351884.dkr.ecr.${AWS::Region}.amazonaws.com/huggingface-pytorch-inference:1.7.1-transformers4.6.1-gpu-py36-cu110-ubuntu18.04
        Mode: SingleModel
        ModelDataUrl: !Sub s3://jumpstart-cache-prod-${AWS::Region}/huggingface-infer/infer-huggingface-summarization-bart-large-cnn-samsum.tar.gz
      ExecutionRoleArn: !GetAtt SageMakerModelExecutionRole.Arn

  LCASMProvisionedSummarizationEndpointConfig:
    Type: AWS::SageMaker::EndpointConfig
    Condition: ProvisionedSageMakerEndpoint
    Properties:
      ProductionVariants:
        - ModelName: !GetAtt LCASMSummarizationModel.ModelName
          InitialInstanceCount: !Ref InitialInstanceCount
          InitialVariantWeight: 1
          InstanceType: !Ref InstanceType
          VariantName: AllTraffic

  LCASMServerlessSummarizationEndpointConfig:
    Type: AWS::SageMaker::EndpointConfig
    Condition: ServerlessSageMakerEndpoint
    Properties:
      ProductionVariants:
        - ModelName: !GetAtt LCASMSummarizationModel.ModelName
          InitialVariantWeight: 1
          VariantName: AllTraffic
          ServerlessConfig:
            MaxConcurrency: 50
            MemorySizeInMB: 4096

  LCASMProvisionedSummarizationEndpoint:
    Type: AWS::SageMaker::Endpoint
    Condition: ProvisionedSageMakerEndpoint
    Properties:
      EndpointConfigName: !GetAtt LCASMProvisionedSummarizationEndpointConfig.EndpointConfigName

  LCASMServerlessSummarizationEndpoint:
    Type: AWS::SageMaker::Endpoint
    Condition: ServerlessSageMakerEndpoint
    Properties:
      EndpointConfigName: !GetAtt 'LCASMServerlessSummarizationEndpointConfig.EndpointConfigName'

  SageMakerModelExecutionRole:
    Type: AWS::IAM::Role
    Properties:
      AssumeRolePolicyDocument:
        Statement:
        - Action:
          - sts:AssumeRole
          Effect: Allow
          Principal:
            Service:
            - sagemaker.amazonaws.com
        Version: '2012-10-17'
      Path: "/"
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/AmazonSageMakerFullAccess

  # Lambda Function to call the summarizer. 
  SummaryLambda:
    Type: AWS::Lambda::Function
    Properties:
      Description: This function summarizes a string from the events object called 'inputs'
      Handler: index.lambda_handler
      Runtime: python3.8
      Role: !GetAtt 'LambdaRole.Arn'
      Timeout: 240
      Environment:
        Variables:
          ENDPOINT_NAME: !If
            - ProvisionedSageMakerEndpoint
            - !GetAtt LCASMProvisionedSummarizationEndpoint.EndpointName
            - !GetAtt LCASMServerlessSummarizationEndpoint.EndpointName 
          FETCH_TRANSCRIPT_LAMBDA_ARN: !Ref FetchTranscriptArn
          PROCESS_TRANSCRIPT: "True"
          TOKEN_COUNT: "1024"

      Code:
        ZipFile: |
          import os
          import io
          import boto3
          from boto3.dynamodb.conditions import Key, Attr
          from botocore.exceptions import ClientError
          import json
          import csv
          import logging
          import re
          
          # grab environment variables
          ENDPOINT_NAME = os.environ['ENDPOINT_NAME']
          FETCH_TRANSCRIPT_LAMBDA_ARN = os.environ['FETCH_TRANSCRIPT_LAMBDA_ARN']
          PROCESS_TRANSCRIPT = (os.getenv('PROCESS_TRANSCRIPT', 'False') == 'True')
          TOKEN_COUNT = int(os.getenv('TOKEN_COUNT', '0'))
          
          runtime = boto3.client('runtime.sagemaker')
          logger = logging.getLogger(__name__)
          ddb = boto3.resource('dynamodb')
          lambda_client = boto3.client('lambda')
          
          html_remover = re.compile('<[^>]*>')
          filler_remover = re.compile('([Uu]m|[Uu]h|[Ll]ike|[Mm]hm)[,]?')
          
          def get_transcripts(callId):
              payload = {
                  'CallId': callId, 
                  'ProcessTranscript': PROCESS_TRANSCRIPT, 
                  'TokenCount': TOKEN_COUNT 
              }
              response = lambda_client.invoke(
                  FunctionName=FETCH_TRANSCRIPT_LAMBDA_ARN,
                  InvocationType='RequestResponse',
                  Payload=json.dumps(payload)
              )
              return response
          
          def lambda_handler(event, context):
              print("Received event: " + json.dumps(event, indent=2))
          
              # Setup model input data using text (utterances) received from LCA
              data = json.loads(json.dumps(event))
              callId = data['CallId']
              transcript_response = get_transcripts(callId)
              transcript_data = transcript_response['Payload'].read().decode()
              transcript_json = json.loads(transcript_data)
              # print(transcript_json)
              payload = {'inputs': transcript_json['transcript']}
          
              summaryText = ""
              try:
                  data = json.dumps(event)
                  response = runtime.invoke_endpoint(EndpointName=ENDPOINT_NAME,
                                                    ContentType='application/json',
                                                    Body=bytes(json.dumps(payload), 'utf-8'))
                  # print(response)
                  result = json.loads(response['Body'].read().decode())
          
                  if len(result) > 0:
                      summary = result[0]
                      # print("Summary: " + summary["generated_text"])
                      summaryText = summary["generated_text"]
                  else:
                      print("No Summary")
                      summaryText = "No summary"
          
              except Exception as e:
                  print("An exception occurred:", e)
                  summaryText = ""
          
              return {"summary": summaryText}
        
        
  
  LambdaRole:
    Type: AWS::IAM::Role
    Properties:
      Description: Summary Lambda Role
      AssumeRolePolicyDocument:
        Version: '2012-10-17'
        Statement:
        - Effect: Allow
          Principal:
            Service: lambda.amazonaws.com
          Action:
          - 'sts:AssumeRole'
      ManagedPolicyArns:
        - arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole
      Policies:
        - PolicyName: !Sub ${AWS::StackName}-InvokeSummarizer
          PolicyDocument:
            Version: "2012-10-17"
            Statement:
              - Effect: "Allow"
                Action: "sagemaker:InvokeEndpoint"
                Resource: !If
                  - ProvisionedSageMakerEndpoint
                  - !Ref LCASMProvisionedSummarizationEndpoint
                  - !Ref LCASMServerlessSummarizationEndpoint
              - Effect: "Allow"
                Action: "dynamodb:Query"
                Resource: !Sub "arn:${AWS::Partition}:dynamodb:${AWS::Region}:${AWS::AccountId}:table/${CallEventsTable}"
              - Effect: "Allow"
                Action: "lambda:InvokeFunction"
                Resource: !Ref FetchTranscriptArn