
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)

[((464, 511), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['PRETRAINED_MODEL'], {}), '(PRETRAINED_MODEL)\n', (493, 511), False, 'from transformers import AutoTokenizer\n'), ((631, 704), 'torch.load', 'torch.load', (['"""../trained_models/t5_model_hotpot_supporting_facts_last.pth"""'], {}), "('../trained_models/t5_model_hotpot_supporting_facts_last.pth')\n", (641, 704), False, 'import torch\n'), ((1087, 1135), 'logging.debug', 'logging.debug', (["('Decoded string' + decoded_string)"], {}), "('Decoded string' + decoded_string)\n", (1100, 1135), False, 'import logging\n'), ((1538, 1585), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['PRETRAINED_MODEL'], {}), '(PRETRAINED_MODEL)\n', (1567, 1585), False, 'from transformers import AutoTokenizer\n'), ((1705, 1774), 'torch.load', 'torch.load', (['"""../trained_models/t5_model_hotpot_full_context_last.pth"""'], {}), "('../trained_models/t5_model_hotpot_full_context_last.pth')\n", (1715, 1774), False, 'import torch\n'), ((2157, 2205), 'logging.debug', 'logging.debug', (["('Decoded string' + decoded_string)"], {}), "('Decoded string' + decoded_string)\n", (2170, 2205), False, 'import logging\n'), ((2621, 2668), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['PRETRAINED_MODEL'], {}), '(PRETRAINED_MODEL)\n', (2650, 2668), False, 'from transformers import AutoTokenizer\n'), ((2788, 2863), 'torch.load', 'torch.load', (['"""../trained_models/bart_model_hotpot_supporting_facts_last.pth"""'], {}), "('../trained_models/bart_model_hotpot_supporting_facts_last.pth')\n", (2798, 2863), False, 'import torch\n'), ((3246, 3294), 'logging.debug', 'logging.debug', (["('Decoded string' + decoded_string)"], {}), "('Decoded string' + decoded_string)\n", (3259, 3294), False, 'import logging\n'), ((3710, 3757), 'transformers.AutoTokenizer.from_pretrained', 'AutoTokenizer.from_pretrained', (['PRETRAINED_MODEL'], {}), '(PRETRAINED_MODEL)\n', (3739, 3757), False, 'from transformers import AutoTokenizer\n'), ((3877, 3948), 'torch.load', 'torch.load', (['"""../trained_models/bart_model_hotpot_full_context_last.pth"""'], {}), "('../trained_models/bart_model_hotpot_full_context_last.pth')\n", (3887, 3948), False, 'import torch\n'), ((4331, 4379), 'logging.debug', 'logging.debug', (["('Decoded string' + decoded_string)"], {}), "('Decoded string' + decoded_string)\n", (4344, 4379), False, 'import logging\n'), ((948, 963), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (961, 963), False, 'import torch\n'), ((2018, 2033), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (2031, 2033), False, 'import torch\n'), ((3107, 3122), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (3120, 3122), False, 'import torch\n'), ((4192, 4207), 'torch.no_grad', 'torch.no_grad', ([], {}), '()\n', (4205, 4207), False, 'import torch\n'), ((246, 271), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (269, 271), False, 'import torch\n'), ((1320, 1345), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (1343, 1345), False, 'import torch\n'), ((2392, 2417), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (2415, 2417), False, 'import torch\n'), ((3481, 3506), 'torch.cuda.is_available', 'torch.cuda.is_available', ([], {}), '()\n', (3504, 3506), False, 'import torch\n'), ((5300, 5331), 'trained_gpt_model.get_inference2', 'get_inference2', (['answer', 'context'], {}), '(answer, context)\n', (5314, 5331), False, 'from trained_gpt_model import get_inference2\n')]