Python Keyphrase Extraction (PKE) AWS Lambda Docker Image (1/3)

lambda_function.py

#!/usr/bin/env python
import json
import pke

extractor = pke.unsupervised.TopicRank()


def lambda_handler(event, context):
    body = event["body"] if "body" in event else event
    if isinstance(body, str):
        body = json.loads(event["body"])

    lang = "en"
    count = 10
    
    if isinstance(body, str):
        text = body
    else:
        text = body["text"] if "text" in body else body
        lang = body["lang"] if "lang" in body else lang
        count = body["count"] if "count" in body else count
    
    extractor.load_document(input=text, language=lang)
    extractor.candidate_selection()
    extractor.candidate_weighting()

    keywords = extractor.get_n_best(n=count)
    tags = [keyword[0] for keyword in keywords]
    
    return {
        "statusCode": 200,
        "tags": tags
    }