
import torch
import logging
# Transformer version 4.9.1 - Newer versions may not work.
from transformers import AutoTokenizer
from trained_gpt_model import get_inference2


def t5_supp_inference(review_text):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # CPU may not work, got to check.
    # device = torch.device('cpu')
    print('Using device:' + str(device))
    PRETRAINED_MODEL = 't5-base'
    SEQ_LENGTH = 600
    tokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL)

    tokenizer.add_special_tokens(
        {'additional_special_tokens': ['<answer>', '<context>']}
    )

    model = torch.load("../trained_models/t5_model_hotpot_supporting_facts_last.pth")
    model.eval()
    encoded_text = tokenizer(
        review_text,
        padding=True,
        max_length=SEQ_LENGTH,
        truncation=True,
        return_tensors="pt"
    ).to(device)

    input_ids = encoded_text['input_ids']
    with torch.no_grad():
        output = model.generate(input_ids)
    decoded_string = tokenizer.decode(output[0], skip_special_tokens=True)
    logging.debug("Decoded string" + decoded_string)
    print(decoded_string)
    # device.empty_cache()
    del model
    del tokenizer
    return decoded_string

def t5_full_inference(review_text):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # CPU may not work, got to check.
    # device = torch.device('cpu')
    print('Using device:' + str(device))
    PRETRAINED_MODEL = 't5-base'
    SEQ_LENGTH = 600
    tokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL)

    tokenizer.add_special_tokens(
        {'additional_special_tokens': ['<answer>', '<context>']}
    )

    model = torch.load("../trained_models/t5_model_hotpot_full_context_last.pth")
    model.eval()
    encoded_text = tokenizer(
        review_text,
        padding=True,
        max_length=SEQ_LENGTH,
        truncation=True,
        return_tensors="pt"
    ).to(device)

    input_ids = encoded_text['input_ids']
    with torch.no_grad():
        output = model.generate(input_ids)
    decoded_string = tokenizer.decode(output[0], skip_special_tokens=True)
    logging.debug("Decoded string" + decoded_string)
    print(decoded_string)
    # device.empty_cache()
    del model
    del tokenizer
    return decoded_string

def bart_supp_inference(review_text):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # CPU may not work, got to check.
    # device = torch.device('cpu')
    print('Using device:' + str(device))
    PRETRAINED_MODEL = 'facebook/bart-base'
    SEQ_LENGTH = 600
    tokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL)

    tokenizer.add_special_tokens(
        {'additional_special_tokens': ['<answer>', '<context>']}
    )

    model = torch.load("../trained_models/bart_model_hotpot_supporting_facts_last.pth")
    model.eval()
    encoded_text = tokenizer(
        review_text,
        padding=True,
        max_length=SEQ_LENGTH,
        truncation=True,
        return_tensors="pt"
    ).to(device)

    input_ids = encoded_text['input_ids']
    with torch.no_grad():
        output = model.generate(input_ids)
    decoded_string = tokenizer.decode(output[0], skip_special_tokens=True)
    logging.debug("Decoded string" + decoded_string)
    print(decoded_string)
    # device.empty_cache()
    del model
    del tokenizer
    return decoded_string

def bart_full_inference(review_text):
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # CPU may not work, got to check.
    # device = torch.device('cpu')
    print('Using device:' + str(device))
    PRETRAINED_MODEL = 'facebook/bart-base'
    SEQ_LENGTH = 600
    tokenizer = AutoTokenizer.from_pretrained(PRETRAINED_MODEL)

    tokenizer.add_special_tokens(
        {'additional_special_tokens': ['<answer>', '<context>']}
    )

    model = torch.load("../trained_models/bart_model_hotpot_full_context_last.pth")
    model.eval()
    encoded_text = tokenizer(
        review_text,
        padding=True,
        max_length=SEQ_LENGTH,
        truncation=True,
        return_tensors="pt"
    ).to(device)

    input_ids = encoded_text['input_ids']
    with torch.no_grad():
        output = model.generate(input_ids)
    decoded_string = tokenizer.decode(output[0], skip_special_tokens=True)
    logging.debug("Decoded string" + decoded_string)
    print(decoded_string)
    # device.empty_cache()
    del model
    del tokenizer
    return decoded_string

# if __name__ == "__main__":
#     review_text = "<answer> a fusional language <context> Typologically, Estonian represents a transitional form from an agglutinating language to a fusional language. The canonical word order is SVO (subject–verb–object)."
#     t5_supp_inference(review_text, md2, device)


def get_inference(answer, context, model_name):
    valuation_text = "<answer> " + answer + " <context> " + context

    if model_name == 't5_supp':
        return t5_supp_inference(valuation_text)
    elif model_name == 't5_full':
        return t5_full_inference(valuation_text)
    elif model_name == 'bart_supp':
        return bart_supp_inference(valuation_text)
    elif model_name == 'bart_full':
        return bart_full_inference(valuation_text)
    elif model_name == 'gpt2':
        return get_inference2(answer, context)

[((15, 16, 15, 63), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', ({(15, 46, 15, 62): 'PRETRAINED_MODEL'}, {}), '(PRETRAINED_MODEL)', False, 'from transformers import AutoTokenizer\n'), ((21, 12, 21, 85), 'torch.load', 'torch.load', ({(21, 23, 21, 84): '"""../trained_models/t5_model_hotpot_supporting_facts_last.pth"""'}, {}), "('../trained_models/t5_model_hotpot_supporting_facts_last.pth')", False, 'import torch\n'), ((35, 4, 35, 52), 'logging.debug', 'logging.debug', ({(35, 18, 35, 51): "('Decoded string' + decoded_string)"}, {}), "('Decoded string' + decoded_string)", False, 'import logging\n'), ((48, 16, 48, 63), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', ({(48, 46, 48, 62): 'PRETRAINED_MODEL'}, {}), '(PRETRAINED_MODEL)', False, 'from transformers import AutoTokenizer\n'), ((54, 12, 54, 81), 'torch.load', 'torch.load', ({(54, 23, 54, 80): '"""../trained_models/t5_model_hotpot_full_context_last.pth"""'}, {}), "('../trained_models/t5_model_hotpot_full_context_last.pth')", False, 'import torch\n'), ((68, 4, 68, 52), 'logging.debug', 'logging.debug', ({(68, 18, 68, 51): "('Decoded string' + decoded_string)"}, {}), "('Decoded string' + decoded_string)", False, 'import logging\n'), ((81, 16, 81, 63), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', ({(81, 46, 81, 62): 'PRETRAINED_MODEL'}, {}), '(PRETRAINED_MODEL)', False, 'from transformers import AutoTokenizer\n'), ((87, 12, 87, 87), 'torch.load', 'torch.load', ({(87, 23, 87, 86): '"""../trained_models/bart_model_hotpot_supporting_facts_last.pth"""'}, {}), "('../trained_models/bart_model_hotpot_supporting_facts_last.pth')", False, 'import torch\n'), ((101, 4, 101, 52), 'logging.debug', 'logging.debug', ({(101, 18, 101, 51): "('Decoded string' + decoded_string)"}, {}), "('Decoded string' + decoded_string)", False, 'import logging\n'), ((114, 16, 114, 63), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', ({(114, 46, 114, 62): 'PRETRAINED_MODEL'}, {}), '(PRETRAINED_MODEL)', False, 'from transformers import AutoTokenizer\n'), ((120, 12, 120, 83), 'torch.load', 'torch.load', ({(120, 23, 120, 82): '"""../trained_models/bart_model_hotpot_full_context_last.pth"""'}, {}), "('../trained_models/bart_model_hotpot_full_context_last.pth')", False, 'import torch\n'), ((134, 4, 134, 52), 'logging.debug', 'logging.debug', ({(134, 18, 134, 51): "('Decoded string' + decoded_string)"}, {}), "('Decoded string' + decoded_string)", False, 'import logging\n'), ((32, 9, 32, 24), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((65, 9, 65, 24), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((98, 9, 98, 24), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((131, 9, 131, 24), 'torch.no_grad', 'torch.no_grad', ({}, {}), '()', False, 'import torch\n'), ((10, 36, 10, 61), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((43, 36, 43, 61), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((76, 36, 76, 61), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((109, 36, 109, 61), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((158, 15, 158, 46), 'trained_gpt_model.get_inference2', 'get_inference2', ({(158, 30, 158, 36): 'answer', (158, 38, 158, 45): 'context'}, {}), '(answer, context)', False, 'from trained_gpt_model import get_inference2\n')]