# -*- coding: utf-8 -*-
"""S22.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1pq0UO46D0emoqF8rPuD4cUznmYVSMESO
"""

# Commented out IPython magic to ensure Python compatibility.
# %pip install lightning -q



import torch
import glob
import math
import sys
import time
from pathlib import Path
from typing import Optional, Tuple, Union

import lightning as L
from lightning.fabric.loggers import CSVLogger
from lightning.fabric.strategies import FSDPStrategy

from tsai_gpt.model import GPT, Block, Config
from tsai_gpt.packed_dataset import CombinedDataset, PackedDataset
from tsai_gpt.speed_monitor import SpeedMonitorBase, estimate_flops, measure_flops
from tsai_gpt.speed_monitor import SpeedMonitorFabric as SpeedMonitor
from tsai_gpt.utils import chunked_cross_entropy, get_default_supported_precision, num_parameters, load_checkpoint
import os
import pickle
from contextlib import nullcontext
from torch.utils.data import DataLoader
import torch.nn.functional as F
from tsai_gpt.tokenizer import Tokenizer
import gradio as gr

model_name = "pythia-160m"
name = "redpajama"
out_dir = Path("out") / name
log_interval = 100

hparams = {k: v for k, v in locals().items() if isinstance(v, (int, float, str)) and not k.startswith("_")}
logger = CSVLogger("out", name, flush_logs_every_n_steps=log_interval)

fabric = L.Fabric(devices=1, strategy='auto', precision=None, loggers=logger)



#print(model.transformer.h[0].mlp.fc.weight)

def generate( model, config, idx, max_new_tokens, temperature=1.0, top_k=None):
        """
        Take a conditioning sequence of indices idx (LongTensor of shape (b,t)) and complete
        the sequence max_new_tokens times, feeding the predictions back into the model each time.
        Most likely you'll want to make sure to be in model.eval() mode of operation for this.

        """
        idx = idx.unsqueeze(dim=0)
        for _ in range(max_new_tokens):
            
            # # if the sequence context is growing too long we must crop it at block_size
            idx_cond = idx if idx.size(1) <= config.block_size else idx[ :,-config.block_size:]
            # forward the model to get the logits for the index in the sequence
            idx_cd = idx
            logits = model(idx_cd)
            # pluck the logits at the final step and scale by desired temperature
            logits = logits[:, -1, :] / temperature
            # optionally crop the logits to only the top k options
            if top_k is not None:
                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < v[:, [-1]]] = -float('Inf')
            # apply softmax to convert logits to (normalized) probabilities
            probs = F.softmax(logits, dim=-1)
            # sample from the distribution
            idx_next = torch.multinomial(probs, num_samples=1)
            # append sampled index to the running sequence and continue
            idx = torch.cat((idx, idx_next), dim=1)

        return idx



checkpoint_dir = Path('./checkpoints/meta-llama/Llama-2-7b-chat-hf')
token = Tokenizer(checkpoint_dir = checkpoint_dir)

def tsaigpt(start:str , max_new_tokens = 300, num_samples =2, tokeniser= token):



  # -----------------------------------------------------------------------------
    temperature = 0.8 # 1.0 = no change, < 1.0 = less random, > 1.0 = more random, in predictions
    top_k = 200 # retain only the top_k most likely tokens, clamp others to have 0 probability
    seed = 1337
    device = 'cpu' # examples: 'cpu', 'cuda', 'cuda:0', 'cuda:1', etc.
    dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 'float32' or 'bfloat16' or 'float16'
    compile = False # use PyTorch 2.0 to compile the model to be faster
    #exec(open('configurator.py').read()) # overrides from command line or config file
    # -----------------------------------------------------------------------------

    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
    torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
    device_type = 'cuda' if 'cuda' in device else 'cpu' # for later use in torch.autocast
    ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
    ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
    checkpoint_path = Path("out/redpajama/iter-031997-ckpt.pth")
    config = Config.from_name(model_name)
    model = GPT(config)

    load_checkpoint(fabric, model, checkpoint_path)
    #print(model)
    model.eval()
    model.to(device)
    if compile:
        model = torch.compile(model) # requires PyTorch 2.0 (optional)



    start_ids = tokeniser.encode(start).to(device)
    #x = torch.tensor(start_ids, dtype=torch.long, device=device).clone().detach()

    # run generation
    with torch.no_grad():
        with ctx:

                y = generate(model =model, config =config ,  max_new_tokens = max_new_tokens, idx = start_ids ,temperature=1.0, top_k=None)
                #print(decode(y[0].tolist()))
                output = tokeniser.decode(y[0])
    return output

INTERFACE = gr.Interface(fn=tsaigpt, inputs=[gr.Textbox(label= "Prompt", value= 'We know what we are, but know not what we may be'),
                    gr.Slider(minimum = 300, maximum = 500, value= 300,  label= "Maximum number of tokens to be generated")] ,
                    outputs=gr.Text(label= "Generated Text"), title="TSAI_GPT",
                 description="TSAIGPT is a transformer-based language model with only 0.16 billion parameters, trained on RedPajama 1T Sample.",
                 examples = [['We know what we are, but know not what we may be',300],]
                        ).launch(debug=True)