
import pandas as pd

from sagemaker.pytorch.model import PyTorchPredictor
from sagemaker.deserializers import JSONDeserializer
from sagemaker.serializers import JSONSerializer
from sagemaker.predictor_async import AsyncPredictor
from sagemaker.predictor import Predictor


def call_label_generator(endpoint_key, questions, sentences, threshold = 0.6):
    multi_label_array = [[] for i in range(len(sentences))]
    
    results_array = []
    
    for question in questions:
        data = {"questions": [question],
                     "sentences" : sentences}
        
        predictor = Predictor(endpoint_name=endpoint_key, serializer=JSONSerializer(), deserializer=JSONDeserializer())
        async_predictor = AsyncPredictor(predictor)
        
        
        result = async_predictor.predict_async(data, "s3://ask-t5-stage/async_inference/ip")
        
        
        """below while loop runs until we get our results, or we get some error """
        while True:
            try:
                results = result.get_result()
                break
            except Exception as ex:
                
                if "Inference could still be running" in str(ex):
                    
                    pass
                
                else:
                    return ex
                    
        
        results_array.append(result)
        for e, individual_sentence_result in enumerate(results):
            if individual_sentence_result[1] > threshold:  
                multi_label_array[e].append(1)
            else:
                multi_label_array[e].append(0)
    
    return multi_label_array, questions, sentences, results_array
            

