import cv2
import torch
import torch.optim as optim
import numpy as np
np.set_printoptions(precision=2, suppress=True)

from AirsimEnv import PeopleAvoidDiscreteEnv, PeopleAvoidDiscreteOneDirEnv
# from DQNConv import DQNConv
from ReplayMemory import ReplayMemory
from arguments import get_args

from workflow import WorkFlow, TorchFlow

class LinearEpsilonAnnealingExplorer(object):
    """
    Exploration policy using Linear Epsilon Greedy

    Attributes:
        start (float): start value
        end (float): end value
        steps (int): number of steps between start and end
    """

    def __init__(self, start, end, steps):
        self._start = start
        self._stop = end
        self._steps = steps

        self._step_size = (end - start) / steps

    def __call__(self, num_actions):
        """
        Select a random action out of `num_actions` possibilities.

        Attributes:
            num_actions (int): Number of actions available
        """
        return np.random.choice(num_actions)

    def _epsilon(self, step):
        """ Compute the epsilon parameter according to the specified step

        Attributes:
            step (int)
        """
        if step < 0:
            return 0. # no exploration if the step is negtive
        elif step > self._steps:
            return self._stop
        else:
            return self._step_size * step + self._start

    def is_exploring(self, step):
        """ Commodity method indicating if the agent should explore

        Attributes:
            step (int) : Current step

        Returns:
             bool : True if exploring, False otherwise
        """
        return np.random.rand() < self._epsilon(step)

class Agent(object):
    def __init__(self, qnet, targetqnet, explorer, actionnum):
        self.qnet = qnet
        self.targetqnet = targetqnet
        self.explorer = explorer
        self.actionnum = actionnum

    def state2input(self, state):
        if state[0].shape[1]>1: # omni-dir
            inputTensor0 = (torch.Tensor(state[0][:,0,:]).cuda(),
                            torch.Tensor(state[0][:,1,:]).cuda(),
                            torch.Tensor(state[0][:,2,:]).cuda(),
                            torch.Tensor(state[0][:,3,:]).cuda())
        else:
            inputTensor0 = torch.Tensor(state[0][:,0,:]).cuda()
        inputTensor1 = torch.Tensor(state[1]).cuda()
        return inputTensor0, inputTensor1

    def net_forward(self, net, state): # shared function by qnet and targetqnet
        '''
        state[0]: n x 1/4 x c x h x w -> 1: one-dir, 4: omni-dir
        state[1]: n x k
        '''
        inputTensor0, inputTensor1 = self.state2input(state)
        q_values = net(inputTensor0, inputTensor1)

        return q_values

    def targetqnet_forward(self, state):
        targetnet_qvalue = self.net_forward(self.targetqnet, state)
        return targetnet_qvalue.detach().cpu().numpy()

    def qnet_forward(self, state):
        return self.net_forward(self.qnet, state)

    def get_random_action(self,):
        return self.explorer(self.actionnum)

    def get_action(self, state, step):

        if self.explorer.is_exploring(step):
            action = self.get_random_action()
            q_mean, q_std = None, None
            # print '   random action', action
        else:
            q_values = self.qnet_forward(state)
            q_values = q_values.detach().cpu().squeeze().numpy()
            #ipdb.set_trace()
            q_mean = np.mean(q_values)
            q_std = np.std(q_values)

            action = np.argmax(q_values)
            # print '    action', action

        return action, q_mean, q_std

    def get_action_vis(self, state): # display the intermediate layers
        inputTensor0, inputTensor1 = self.state2input(state)
        q_values, disps = self.qnet.forward_disp(inputTensor0, inputTensor1)
        q_values = q_values.detach().cpu().squeeze().numpy()
        disps_np = []
        for disp in disps:
            disp = disp.detach().cpu().squeeze().numpy()
            disps_np.append(disp)
        # import ipdb;ipdb.set_trace()
        return q_values, disps_np

class QLearning():
    def __init__(self, agent, optimizer, replaymemory, gamma, batchsize):
        self.agent = agent
        self.optimizer = optimizer
        self.memory = replaymemory
        self.gamma = gamma
        self.batchsize = batchsize
        self.criterion = torch.nn.SmoothL1Loss() # for baseline

    def compute_q_targets(self, state, rewards, terminals): # TOTEST
        ret_vals = []
        #ipdb.set_trace()
        post_states_q = self.agent.targetqnet_forward(state)
        #ipdb.set_trace()
        for i, terminal in enumerate(terminals):
            if terminal:
                ret_vals.append(rewards[i])
            else:
                ret_vals.append(rewards[i] + self.gamma*np.max(post_states_q[i]))
        return torch.Tensor(ret_vals)

    def update(self,):
        pre_states, actions, post_states, rewards, terminals = self.memory.minibatch(self.batchsize)

        q_targets = self.compute_q_targets(post_states, rewards, terminals)

        actions = torch.from_numpy(actions).type(torch.LongTensor).cuda()
        q_vals_pre_states =  self.agent.qnet_forward(pre_states).gather(1, actions.view(-1,1))
        loss = self.criterion(q_vals_pre_states.squeeze(-1), q_targets.cuda())

        # Optimize the model
        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        return {'loss':loss.item()}

class TrainDQNPeopleAvoidance(TorchFlow.TorchFlow):
    def __init__(self, workingDir, args, prefix = "", suffix = "", plotterType = 'Visdom'):
        super(TrainDQNPeopleAvoidance, self).__init__(workingDir, prefix, suffix, disableStreamLogger = True, plotterType = plotterType)
        self.args = args
        self.prefix = prefix
        self.saveModelName = 'dqn'
        self.frameNum = args.multi_frame

        if args.omni_dir:
            self.imgshape = (4, 1+self.frameNum, 16, 160)
        else:
            self.imgshape = (1, 1+self.frameNum, 16, 160)
        self.velshape = (15,)

        self.lr = self.args.lr
        self.period = 0.1/self.args.speedup

        self.AV['loss'].avgWidth = 1000
        self.add_accumulated_value('reward',10)
        self.add_accumulated_value('qmean', 1000)
        self.add_accumulated_value('qstd', 1000)
        self.append_plotter("loss", ['loss'], [True])
        self.append_plotter("reward", ['reward'], [True])
        self.append_plotter("qvalue", ['qmean', 'qstd'], [True, True])

        self.countTrain = 0
        self.countTest = 0

        if self.args.omni_dir:
            self.env =  PeopleAvoidDiscreteEnv(period=self.period, framenum=self.frameNum, strictsafe=args.strict_safe) # DummyEnv() #
        else:
            self.env =  PeopleAvoidDiscreteOneDirEnv(period=self.period, strictsafe=args.strict_safe)
        self.current_state = self.env.reset()
        self.current_state = self.stateTransform(self.current_state)
        self.action_num = self.env.action_space.shape[0]
        self.episode_reward = []

        from DQNConv import DQNConv
        self.dqn = DQNConv(self.action_num, self.velshape[0], omni_dir=self.args.omni_dir, input_num=self.frameNum+1)
        self.targetdqn = DQNConv(self.action_num, self.velshape[0], omni_dir=self.args.omni_dir, input_num=self.frameNum+1)

        if self.args.load_qnet:
            dqn_model = self.args.working_dir + '/models/' + self.args.qnet_model
            self.load_model(self.dqn, dqn_model)
        self.targetdqn.load_state_dict(self.dqn.state_dict())
        self.dqn.cuda()
        self.targetdqn.cuda()

        epsilon_start, epsilon_end = float(self.args.epsilon.split('_')[0]), float(self.args.epsilon.split('_')[1])
        self.agent = Agent(self.dqn, self.targetdqn, LinearEpsilonAnnealingExplorer(epsilon_start, epsilon_end, self.args.train_step), self.action_num)
        self.memory = ReplayMemory(self.args.memory_size, self.imgshape, self.velshape)
        self.optimizer = optim.Adam(self.dqn.parameters(), lr = self.lr, eps=1e-5) #,weight_decay=1e-5)
        self.algo = QLearning(self.agent, self.optimizer, self.memory, self.args.gamma, self.args.batch_size)

        if self.args.load_memory:
            self.memory.load_memory(prefix=self.args.memory_prefix)

    def initialize(self):
        super(TrainDQNPeopleAvoidance, self).initialize()
        # fill memory
        if not self.args.test and len(self.memory)<self.args.batch_size:
            for k in range(self.args.batch_size):
                action = self.agent.get_random_action()
                state, reward, finish, _ = self.env.step(action)
                self.current_state = self.stateTransform(state)
                self.episode_reward.append(reward)
                self.memory.append(self.current_state[0][0], action, reward, finish, self.current_state[1][0])
                if finish:
                    self.episode_reward = []

                    self.current_state = self.env.reset()
                    self.current_state = self.stateTransform(self.current_state)

        if self.args.test:
            self.droneimg = self.loadDroneImg()

        super(TrainDQNPeopleAvoidance, self).post_initialize()

    def stateTransform(self, state):
        '''
        state: state['img']: 1/4 x 2 x 80 x 32
               state['vel']: 3
        return: state[0]: 1 x 1/4 x 2 x 80 x 32 --> handle the discripency between sampling (1 x 2 x 80 x 32) and training (batch x 1 x 2 x 80 x 32)
                state[1]: 1 x 30
        '''
        state_img = state['img'].copy()
        state_img[:,0,:,:] = state_img[:,0,:,:].clip(0.2) # minimun distance 0.2m
        state_img[:,0,:,:] = 1.0/state_img[:,0,:,:] # inverse distance

        state_vel = np.tile(state['vel'], 5) # repeat the vel input
        return np.expand_dims(state_img,axis=0), state_vel.reshape((1,)+state_vel.shape)

    def train(self):
        super(TrainDQNPeopleAvoidance, self).train()
        self.countTrain += 1

        action, qmean, qstd = self.agent.get_action(self.current_state, self.countTrain)

        state, reward, finish,_ = self.env.step(action)
        self.current_state = self.stateTransform(state)
        self.episode_reward.append(reward)
        self.memory.append(self.current_state[0][0], action, reward, finish, self.current_state[1][0])

        if qmean is not None and qstd is not None:
            self.AV['qmean'].push_back(qmean, self.countTrain)
            self.AV['qstd'].push_back(qstd, self.countTrain)

        if finish:
            self.AV['reward'].push_back(sum(self.episode_reward), self.countTrain)
            self.episode_reward = []

            self.current_state = self.env.reset()
            self.current_state = self.stateTransform(self.current_state)

        #ipdb.set_trace()
        if self.countTrain % self.args.train_interval == 0:

            # import ipdb;ipdb.set_trace()
            plot_ret = self.algo.update()
            for item in plot_ret.keys():
                self.AV[item].push_back(plot_ret[item], self.countTrain)

        if ( self.countTrain % self.args.plot_interval == 0):
            self.plot_accumulated_values()
            losslogstr = self.get_log_str()
            self.logger.info("%s #%d - %s lr: %.6f"  % (self.args.exp_prefix[:-1],
                self.countTrain, losslogstr, self.lr))

        if ( self.countTrain % self.args.snapshot == 0 ):
            self.write_accumulated_values()
            self.draw_accumulated_values()
            self.save_model(self.dqn, self.saveModelName+'_'+str(self.countTrain))
            self.memory.save_memory(prefix=self.prefix)

        if self.countTrain % self.args.target_update_interval == 0:
            self.targetdqn.load_state_dict(self.dqn.state_dict())
            print('===> Update target_net...')
            print('')

    def test(self, ):

        self.countTest += 1

        action, qmean, qstd = self.agent.get_action(self.current_state, -1)
        state, reward, finish, _ = self.env.step(action)
        self.current_state = self.stateTransform(state)
        self.episode_reward.append(reward)
        self.memory.append(self.current_state[0][0], action, reward, finish, self.current_state[1][0])
        self.visualizeAction(action)

        print('')
        print('{}, action {}'.format(self.countTest, action))
        if finish:
            print('reward for this episode: {}'.format(sum(self.episode_reward)))
            self.episode_reward = []

            self.current_state = self.env.reset()
            self.current_state = self.stateTransform(self.current_state)

            return True

        return False

    def action2coord(self, action, centerpt=(25,25), linelen=20):
        if action==0:
            endpt = centerpt
            return endpt
        # action = (action+5)%8 + 1 # adjust the visulization direction
        if action==1:
            endpt = (centerpt[0], centerpt[1]-linelen)
        elif action==2:
            endpt = (centerpt[0]+linelen, centerpt[1]-linelen)
        elif action==3:
            endpt = (centerpt[0]+linelen, centerpt[1])
        elif action==4:
            endpt = (centerpt[0]+linelen, centerpt[1]+linelen)
        elif action==5:
            endpt = (centerpt[0], centerpt[1]+linelen)
        elif action==6:
            endpt = (centerpt[0]-linelen, centerpt[1]+linelen)
        elif action==7:
            endpt = (centerpt[0]-linelen, centerpt[1])
        elif action==8:
            endpt = (centerpt[0]-linelen, centerpt[1]-linelen)
        else:
            print('action error..')
        return endpt

    def visualizeAction(self, action):
        imgsize = 50
        startpt = (int(imgsize/2),int(imgsize/2))
        endpt = self.action2coord(action,startpt)

        img=np.ones((imgsize,imgsize)).astype(np.uint8)*255
        cv2.imshow('action',cv2.arrowedLine(img,startpt,endpt,color=(0,0,255),thickness=2))
        cv2.waitKey(2)

    def activation2imgs(self, disps):
        activationlist = []
        shapelist = []
        for k,disp in enumerate(disps):
            scale = k/2+1
            disp_norm = (disp/disp.std()).clip(0,1)
            if len(disp_norm.shape)==1: # FC layers
                disp_norm = disp_norm.reshape(disp_norm.shape[-1],1)
                disp_norm = cv2.resize(disp_norm, (0,0), fx=4, fy=4, interpolation=3)
            else:
                if len(disp_norm.shape)==2:
                    disp_norm = np.expand_dims(disp_norm,axis=1)
                # add gray strip
                disp_norm = np.concatenate((disp_norm,np.ones_like(disp_norm)[:,0:int(4/scale+1),:]*0.9),axis=1)
                disp_norm = disp_norm.reshape(-1,disp_norm.shape[-1])
                disp_norm = cv2.resize(disp_norm, (0,0), fx=scale, fy=scale, interpolation=3)
            # imshow
            activationlist.append(disp_norm)
            shapelist.append(disp_norm.shape)
        return activationlist, shapelist

    def state2img(self, state):
        imgshape = state.shape # camnum x frames x h x w
        smallspace = np.ones((10,imgshape[-1]),dtype=np.float32) * 0.9
        # imgh = imgshape[2]*imgshape[0] + 10 * (imgshape[0]-1)
        # imgdisp = np.zeros((imgh, imgshape[3]),dtype=np.uint8)
        imgdisp = np.ones((0, imgshape[-1]), dtype=np.float32)
        for k in range(imgshape[0]):
            depthdisp = np.clip(state[k,0,:,:]*2.0/256.0,0,1)
            imgdisp = np.concatenate((imgdisp,depthdisp,smallspace),axis=0)
            for w in range(1, imgshape[1]):
                imgdisp = np.concatenate((imgdisp,state[k,w,:,:],smallspace),axis=0)
            imgdisp = np.concatenate((imgdisp,smallspace),axis=0)
        # cv2.imshow('input',imgdisp)
        # cv2.waitKey(1)
        return imgdisp

    def qvalue2img(self, qvalues):
        imgsize = 200
        startpt = (int(imgsize/2),int(imgsize/2))
        qimg=np.ones((imgsize,imgsize)).astype(np.float32) * 0.9
        qvis = (qvalues-qvalues.mean())/qvalues.std()
        qvis = np.exp(qvis)*2
        qvis = qvis/qvis.sum()
        for k,q in enumerate(qvis):
            endpt = self.action2coord(k,startpt,linelen=int(imgsize*q))
            cv2.arrowedLine(qimg,startpt,endpt,color=(0,0,1),thickness=2)
        print('Q values: {}'.format(qvalues))
        # cv2.imshow('qvalue',qimg)
        # cv2.waitKey(1)
        return qimg

    def loadDroneImg(self, imgname='droneimg.png'):
        droneimg = cv2.imread(imgname).astype(np.float32)
        droneimg = np.clip(droneimg/255.0, 0, 1)
        droneimg = cv2.resize(droneimg, (0,0), fx=0.3, fy=0.3)
        return droneimg

    def alignCenterImgs(self, imglist, shapelist, imggap = 50):
        imgw, imgh = 0, 0
        for imgshape in shapelist:
            if imgh<imgshape[0]:
                imgh = imgshape[0]
            imgw += (imggap + imgshape[1])
        imgh += imggap
        alignimg = np.ones((imgh,imgw,3),dtype=np.float32) * 0.9
        startx, starty = 0, int(imggap/2)
        for imgshape,img in zip(shapelist,imglist):
            startx = int(imgh/2-imgshape[0]/2) # align middle
            if len(imgshape)==2:
                img = np.tile(np.expand_dims(img,axis=2),3)
            alignimg[startx:startx+imgshape[0],starty:starty+imgshape[1],:] = img
            starty += (imggap + imgshape[1])
        return alignimg

    def vis_network(self, state, qvalue, disps):
        activationlist, shapelist = self.activation2imgs(disps)
        inputimg = self.state2img(state['img'])
        qimg = self.qvalue2img(qvalue)
        # draw images together
        actimg = self.alignCenterImgs(activationlist, shapelist, imggap = 30)
        dispimg = self.alignCenterImgs([inputimg,actimg],[inputimg.shape, actimg.shape],imggap=50)
        dispimg = self.alignCenterImgs([dispimg,self.droneimg,qimg],[dispimg.shape, self.droneimg.shape, qimg.shape],imggap=10)

        cv2.imshow('activation',dispimg)
        cv2.waitKey(1)

    def test_vis(self, ):

        self.countTest += 1

        qvalue, disps = self.agent.get_action_vis(self.current_state)

        action = np.argmax(qvalue)
        state, reward, finish, _ = self.env.step(action)

        # self.env.render(state)
        self.current_state = self.stateTransform(state)
        self.episode_reward.append(reward)
        # self.memory.append(self.current_state[0][0], action, reward, finish, self.current_state[1][0])
        self.visualizeAction(action)
        self.vis_network(state, qvalue, disps)

        print('')
        print('{}, action {}'.format(self.countTest, action))
        if finish:
            print('reward for this episode: {}'.format(sum(self.episode_reward)))
            self.episode_reward = []

            self.current_state = self.env.reset()
            self.current_state = self.stateTransform(self.current_state)

            return True

        return False

    def finalize(self):
        super(TrainDQNPeopleAvoidance, self).finalize()
        if not self.args.test and self.countTrain != self.args.train_step:
            self.write_accumulated_values()
            self.draw_accumulated_values()
            self.save_model(self.dqn, self.saveModelName+'_'+str(self.countTrain))
            self.memory.save_memory(prefix=self.prefix)
        self.env.close()
        self.logger.info("Finalized.")

if __name__ == '__main__':

    args = get_args()
    if args.use_int_plotter:
        plottertype = 'Int'
    else:
        plottertype = 'Visdom'
    try:
        # Instantiate an object for MyWF.
        traindqn = TrainDQNPeopleAvoidance(args.working_dir, args, prefix = args.exp_prefix, plotterType = plottertype)

        # Initialization.
        traindqn.initialize()
        while True:
            if not args.test:
                traindqn.train()
                if (traindqn.countTrain>=args.train_step):
                    break
            else:
                if traindqn.test_vis():
                    break

        traindqn.finalize()

    except WorkFlow.SigIntException as sie:
        print( sie.describe() )
        print( "Quit after finalize." )
        traindqn.finalize()
    except WorkFlow.WFException as e:
        print( e.describe() )
        traindqn.finalize()

    print("Done.")


#!/usr/bin/env python3
import json
import random
import logging
import operator
import numpy as np

from sqlalchemy import Float
from sqlalchemy import Column
from sqlalchemy import String
from sqlalchemy import Integer
from sqlalchemy import ForeignKey

from gs.dataserver.model import Base
from gs.dataserver.model import _batch
from gs.dataserver.model import ModelDB

from gs.dataserver.model.product import Product
from gs.dataserver.model.product import PRODUCT_ID_LENGTH

from gs.dataserver.model.user import ReviewUser
from gs.dataserver.model.user import get_user_set
from gs.dataserver.model.user import USER_ID_LENGTH

from gs.dataserver.algorithms import GraphSVD

class Review(Base):
    """
    ORM object representing a Review Entity
    """

    __tablename__ = 'review'

    reviewerID = Column(String(USER_ID_LENGTH), ForeignKey("reviewuser.id"), primary_key=True)
    asin = Column(String(PRODUCT_ID_LENGTH), ForeignKey("product.asin"), primary_key=True)
    helpful = Column(String)
    text = Column(String)
    overall = Column(Float)
    summary = Column(String)
    time = Column(Integer)

    _JSON_ATTS = [
        "helpful"
    ]

    def __init__(self, *args, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)

    def convert_json(self):
        """
        Convert the object to JSON
        """
        for att in self._JSON_ATTS:
            val = getattr(self, att, "[]") or "[]"
            setattr(self, att, json.loads(val))

class Recommendation(Base):
    """
    ORM object for a Recommendation entity
    """

    __tablename__ = "recommendation"

    userID = Column(String(USER_ID_LENGTH), ForeignKey("reviewuser.id"), primary_key=True)
    asin = Column(String(PRODUCT_ID_LENGTH), ForeignKey("product.asin"), primary_key=True)
    estimated_rating = Column(Float)

    def __init__(self, *args, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)

def reviewfunc(func, *args, **kwargs):
    """
    Check for the review table, create if
    not present
    """
    def wrapper(*args, **kwargs):
        with ModelDB() as db:
            if "review" not in db.inspector.get_table_names():
                Base.metadata.create_all(db._engine)
        return func(*args, **kwargs)
    return wrapper

def recommendationfunc(func, *args, **kwargs):
    """
    Check for recommendation table, create if not present
    """
    def wrapper(*args, **kwargs):
        with ModelDB() as db:
            if "recommendation" not in db.inspector.get_table_names(): # pragma: no cover
                Base.metadata.create_all(db._engine)
        return func(*args, **kwargs)
    return wrapper

@reviewfunc
def add_reviews(objs, batch_size=50):
    """
    Add a set of reviews
    """
    with ModelDB() as db:
        for b in _batch(objs, batch_size):
            db.session.bulk_save_objects(b)

@reviewfunc
def get_reviews(rids=None, pids=None):
    """
    Get all reveiws for the given pids and rids.
    """
    with ModelDB() as db:
        query = db.session.query(Review)
        if rids is not None:
            query = query.filter(Review.reviewerID.in_(rids))
        if pids is not None:
            query = query.filter(Review.asin.in_(pids))
        result = query.all()
    for r in result:
        r.convert_json()
    return result

@reviewfunc
def review_exists(rid, pid):
    """
    Check if a review exists.
    """
    with ModelDB() as db:
        query = db.session.query(Review).filter(Review.reviewerID == rid).filter(Review.asin == pid)
        ans = db.session.query(query.exists()).all().pop()[0]
    return ans

@recommendationfunc
def add_recommendations(objs, batch_size=50):
    """
    Add a set of recommendations.
    """
    with ModelDB() as db:
        for b in _batch(objs, batch_size):
            db.session.bulk_save_objects(b)

@recommendationfunc
def get_recommendations(uids=None):
    """
    Get recommendations for a given set of users.
    """
    with ModelDB() as db:
        query = db.session.query(Recommendation)
        if uids is not None:
            query = query.filter(Recommendation.userID.in_(uids))
        result = query.all()
    return result

REC_LIMIT = 10000

def compute_recommendations(components=3):
    """
    Compute the new recommendations.
    """
    logging.info('Collecting data...')
    # Get the users and corresponding and reviews
    users = get_user_set(REC_LIMIT)
    reviews = get_reviews(rids=users)
    pids = list(set([r.asin for r in reviews]))
    logging.info("A total of {} users and {} products are being used.".format(len(users), len(pids)))
    mat = []
    holdout = []
    logging.info('Building matrix...')
    for u in users:
        mat.append([0 for _ in range(len(pids))])
        # Get all relevant reviews
        user_reviews = [r for r in reviews if r.reviewerID == u]
        for r in user_reviews:
            # Flip a coin, if heads save for testing
            if len(holdout) < 200 and random.randint(1, 2) == 1:
                holdout.append((r.reviewerID, r.asin, r.overall / 5.0))
            # insert values into matrix
            idx = pids.index(r.asin)
            mat[-1][idx] = r.overall / 5.0
    logging.info('Computing SVD...')
    # Compute the SVD object
    recommender = GraphSVD(np.array(mat), components=components, epsilon=1e-6)
    logging.info('Getting test results...')
    correct = 0
    # for every item held out, test against
    # the predictions
    for _id, asin, score in holdout:
        idx = users.index(_id)
        pidx = pids.index(asin)
        predicted = recommender.predictions[idx][pidx]
        if abs(predicted - score) <= 0.2:
                correct += 1
    # Log success rate
    logging.info('Success rate: {}%.'.format((correct / 200)*100))
    logging.info('Uploading recommendations...')
    # Upload all the recommendatinos
    for idx, u in enumerate(users):
        # Sort based on likliehood
        recommended = sorted(zip(pids, recommender.predictions[idx]), key=operator.itemgetter(1), reverse=True)
        recommended = recommended[:5]
        items = [Recommendation(userID=u, asin=r[0], estimated_rating=r[1]) for r in recommended]
        logging.info('Uploading recommendations for user {}...'.format(u))
        add_recommendations(items)


import math
from abc import abstractmethod
from pathlib import Path
from typing import List, Literal, Optional, Union
import os

import numpy as np

from . import comm
from . import profiler
from .gpt import GptInitModelParameters

from nebullvm.optional_modules.torch import torch

PathLike = Union[str, Path]

def to_numpy_dtype(maybe_str_dtype: Union[str, np.dtype]):
    assert isinstance(maybe_str_dtype, (str, np.dtype))
    if isinstance(maybe_str_dtype, str):
        try:
            dtype = {
                "fp16": np.float16,
                "float16": np.float16,
                "fp32": np.float32,
                "float32": np.float32,
            }[maybe_str_dtype]
        except KeyError:
            raise ValueError(
                f"Cannot convert to numpy data type, got {maybe_str_dtype}"
            )
    else:
        dtype = maybe_str_dtype
    return dtype

def to_torch_dtype(maybe_str_dtype: Union[str, torch.dtype]):

    if isinstance(maybe_str_dtype, torch.dtype):
        dtype = maybe_str_dtype
    else:
        try:
            dtype = {
                "bf16": torch.bfloat16,
                "fp16": torch.float16,
                "fp32": torch.float32,
                "bfloat16": torch.bfloat16,
                "float16": torch.float16,
                "float32": torch.float32,
            }[maybe_str_dtype]
        except KeyError:
            raise ValueError(
                f"Cannot convert to torch data type, got {maybe_str_dtype}"
            )
    return dtype

def load_weight_from_bin(
    checkpoint_path: PathLike,
    shape: List[int],
    weight_dtype: Union[str, np.dtype],
):
    """Load a weight from a bin file.

    # Args.
        checkpoint_path: str or Path,
            a checkpoint file path of an FT's layer weight.
        shape: list of int, the shape of weight tensor.
        weight_dtype: str or np.dtype, the data type of the stored weight.
    """
    weight_dtype = to_numpy_dtype(weight_dtype)
    return torch.from_numpy(np.fromfile(checkpoint_path, dtype=weight_dtype))

LayernormType = Literal["pre_layernorm", "post_layernorm"]

class GptLayerWeights:
    def __init__(
        self,
        num_heads: int,
        size_per_head: int,
        inter_size: int,
        num_layers: int,
        tensor_para_size: int = 1,
        pipeline_para_size: int = 1,
        has_adapters: bool = False,
        adapter_inter_size: int = 0,
        int8_mode: int = 0,
    ):

        assert num_heads % tensor_para_size == 0, (
            f"num_heads ({num_heads}) is not multiple of "
            "tensor para size ({tensor_para_size})"
        )

        self.num_heads = num_heads
        self.size_per_head = size_per_head
        self.hidden_units = num_heads * size_per_head
        self.num_layers = num_layers

        self.tensor_para_size = tensor_para_size
        self.tensor_para_rank = comm.get_tensor_para_rank()
        self.pipeline_para_size = pipeline_para_size
        self.pipeline_para_rank = comm.get_pipeline_para_rank()

        self.has_adapters = has_adapters
        self.adapter_inter_size = adapter_inter_size

        self.local_num_layers = num_layers // pipeline_para_size
        self.local_num_heads = num_heads // tensor_para_size
        self.local_hidden_units = self.local_num_heads * size_per_head
        self.local_inter_size = inter_size // tensor_para_size
        self.local_adapter_inter_size = (
            self.adapter_inter_size // tensor_para_size
        )

        self.weight_transpose_calibrate_quantize = None
        assert int8_mode in [0, 1], "Invalid int8 mode for GPT. Must be 0 or 1"
        self.int8_mode = int8_mode
        if self.int8_mode == 1:
            quant = (
                torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix  # noqa: E501
            )
            self.weight_transpose_calibrate_quantize = lambda x: quant(
                x, torch.int8
            )

        self.weights = None
        self.int8_weights = None
        self.int8_scales = None

        self.expected_weight_shapes = list()

        # pylint:disable=line-too-long
        # Transformer blocks
        self.expected_weight_shapes.extend(
            [(self.hidden_units,)] * self.local_num_layers
        )  # input layernorm weight
        self.expected_weight_shapes.extend(
            [(self.hidden_units,)] * self.local_num_layers
        )  # input layernorm bias
        self.expected_weight_shapes.extend(
            [(self.hidden_units, self.local_hidden_units * 3)]
            * self.local_num_layers
        )  # attention qkv weight
        self.expected_weight_shapes.extend(
            [(self.local_hidden_units * 3,)] * self.local_num_layers
        )  # attention qkv bias
        self.expected_weight_shapes.extend(
            [(self.local_hidden_units, self.hidden_units)]
            * self.local_num_layers
        )  # attention dense weight
        self.expected_weight_shapes.extend(
            [(self.hidden_units,)] * self.local_num_layers
        )  # attention dense bias
        self.expected_weight_shapes.extend(
            [(self.hidden_units,)] * self.local_num_layers
        )  # post attention layernorm weight
        self.expected_weight_shapes.extend(
            [(self.hidden_units,)] * self.local_num_layers
        )  # post attention layernorm bias
        self.expected_weight_shapes.extend(
            [(self.hidden_units, self.local_inter_size)]
            * self.local_num_layers
        )  # ffn_kernel1
        self.expected_weight_shapes.extend(
            [(self.local_inter_size,)] * self.local_num_layers
        )  # ffn_bias1
        self.expected_weight_shapes.extend(
            [(self.local_inter_size, self.hidden_units)]
            * self.local_num_layers
        )  # ffn_kernel2
        self.expected_weight_shapes.extend(
            [(self.hidden_units,)] * self.local_num_layers
        )  # ffn_bias2

        # Adapters
        if self.has_adapters:
            self.expected_weight_shapes.extend(
                [(self.hidden_units, self.local_adapter_inter_size)]
                * self.local_num_layers
            )  # adaptor1_kernel1
            self.expected_weight_shapes.extend(
                [(self.local_adapter_inter_size,)] * self.local_num_layers
            )  # adaptor1_bias1
            self.expected_weight_shapes.extend(
                [(self.local_adapter_inter_size, self.hidden_units)]
                * self.local_num_layers
            )  # adaptor1_kernel2
            self.expected_weight_shapes.extend(
                [(self.hidden_units,)] * self.local_num_layers
            )  # adaptor1_bias2
            self.expected_weight_shapes.extend(
                [(self.hidden_units, self.local_adapter_inter_size)]
                * self.local_num_layers
            )  # adaptor2_kernel1
            self.expected_weight_shapes.extend(
                [(self.local_adapter_inter_size,)] * self.local_num_layers
            )  # adaptor2_bias1
            self.expected_weight_shapes.extend(
                [(self.local_adapter_inter_size, self.hidden_units)]
                * self.local_num_layers
            )  # adaptor2_kernel2
            self.expected_weight_shapes.extend(
                [(self.hidden_units,)] * self.local_num_layers
            )  # adaptor2_bias2
        # pylint:enable=line-too-long

    @classmethod
    def from_config(cls, config: GptInitModelParameters):
        return cls(
            num_heads=config.head_num,
            size_per_head=config.size_per_head,
            inter_size=4 * config.head_num * config.size_per_head,
            num_layers=config.layer_num,
            tensor_para_size=config.tensor_para_size,
            pipeline_para_size=config.pipeline_para_size,
            has_adapters=config.has_adapters,
            adapter_inter_size=config.adapter_inter_size,
            int8_mode=config.int8_mode,
        )

    @property
    def dtype(self):
        return self.weights[0].dtype

    @property
    def device(self):
        return self.weights[0].device

    def _map(self, func):
        for i in range(len(self.weights)):
            if isinstance(self.weights[i], list):
                for j in range(len(self.weights[i])):
                    self.weights[i][j] = func(self.weights[i][j])
            else:
                self.weights[i] = func(self.weights[i])

    def _map_int8(self, func):
        for i in range(len(self.int8_weights)):
            if isinstance(self.int8_weights[i], list):
                for j in range(len(self.int8_weights[i])):
                    self.int8_weights[i][j] = func(self.int8_weights[i][j])

            else:
                self.int8_weights[i] = func(self.int8_weights[i])
        for i in range(len(self.int8_scales)):
            if isinstance(self.int8_scales[i], list):
                for j in range(len(self.int8_scales[i])):
                    self.int8_scales[i][j] = func(self.int8_scales[i][j])
            else:
                self.int8_scales[i] = func(self.int8_scales[i])

    def float(self):
        if self.dtype == torch.float32:
            return
        self._map(lambda x: x.float())

    def half(self):
        if self.dtype == torch.float16:
            return
        self._map(lambda x: x.half())
        if self.int8_mode == 1:
            self._map_int8(lambda w: w.half())

    def bfloat16(self):
        if self.dtype == torch.bfloat16:
            return
        self._map(lambda x: x.bfloat16())
        if self.int8_mode == 1:
            self._map_int8(lambda w: w.bfloat16())

    def cuda(self, device=None):
        self._map(lambda x: x.cuda(device))
        if self.int8_mode == 1:
            self._map_int8(lambda x: x.cuda(device))

    def to(self, device=None):
        self._map(lambda x: x.to(device))
        if self.int8_mode == 1:
            self._map_int8(lambda x: x.to(device))

    def is_valid_pp_group(self, layer, pp_rank):
        return layer // self.layers_per_device == pp_rank

    def load(
        self,
        checkpoint_path: PathLike,
        compute_dtype: torch.dtype,
        weight_dtype: Optional[Union[str, np.dtype]] = None,
        device: Optional[Union[int, str, torch.device]] = None,
    ):
        """Load checkpoint weights.

        # Args.
            checkpoint_path: str or Path,
                a checkpoint directory where FT checkpoint files locate.
            weight_dtype: str or np.dtype, the data type of stored weights.
        """

        checkpoint_path = Path(checkpoint_path)
        if not checkpoint_path.exists():
            raise FileNotFoundError(
                f"Could not find checkpoint {str(checkpoint_path)}"
            )

        weight_dtype = to_numpy_dtype(weight_dtype)
        print(
            f"Load weights from {str(checkpoint_path)} (data type: {weight_dtype}"  # noqa: E501
        )

        self.weights = list()
        self.int8_weights = list()
        self.int8_scales = list()
        torch.cuda.empty_cache()

        def _load_from_file(fname):
            quant_sub_names = [
                "attention.query_key_value.weight",
                "attention.dense.weight",
                "dense_h_to_4h.weight",
                "dense_4h_to_h.weight",
            ]
            _weight = torch.from_numpy(
                np.fromfile(checkpoint_path / fname, dtype=weight_dtype)
            )
            _weight = _weight.to(compute_dtype)
            weight_index = len(self.weights)
            expected_shape = self.expected_weight_shapes[weight_index]

            try:
                if _weight.nelement() > 0:
                    _weight = _weight.reshape(expected_shape)
            except:  # noqa: E722
                raise ValueError(
                    f"num_heads, size_per_head, vocab_size, and max_seq_len must be the same "  # noqa: E501
                    f"as the ones during training (weight: {fname} expected shape: {expected_shape}, "  # noqa: E501
                    f"got shape: {_weight.shape})."
                )

            should_quantize = any(
                sub_name in fname for sub_name in quant_sub_names
            )
            if self.int8_mode != 0 and should_quantize:
                calibrate = self.weight_transpose_calibrate_quantize
                int8_weight, int8_scales = calibrate(_weight)

                # int8 weights should appear in same order as FP weights.
                # Move to device and add to the int8 list.
                dummy_weight = torch.empty(0, dtype=compute_dtype)
                if device is not None:
                    int8_weight = int8_weight.to(device)
                    int8_scales = int8_scales.to(device)
                    dummy_weight = dummy_weight.to(device)

                self.int8_weights.append(int8_weight)
                self.int8_scales.append(int8_scales)
                self.weights.append(dummy_weight)
            else:
                if device is not None:
                    _weight = _weight.to(device)
                self.weights.append(_weight)

        # Load
        # pylint:disable=line-too-long
        layer_offset = self.local_num_layers * self.pipeline_para_rank
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.input_layernorm.weight.bin"
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.input_layernorm.bias.bin"
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.attention.query_key_value.weight.{self.tensor_para_rank}.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.attention.query_key_value.bias.{self.tensor_para_rank}.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.attention.dense.weight.{self.tensor_para_rank}.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.attention.dense.bias.bin"
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.post_attention_layernorm.weight.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.post_attention_layernorm.bias.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.mlp.dense_h_to_4h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.mlp.dense_h_to_4h.bias.{self.tensor_para_rank}.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.mlp.dense_4h_to_h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
            )
            for i in range(self.local_num_layers)
        ]
        [
            _load_from_file(
                f"model.layers.{layer_offset + i}.mlp.dense_4h_to_h.bias.bin"
            )
            for i in range(self.local_num_layers)
        ]

        if self.has_adapters:
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_attention_adapter.dense_h_to_4h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_attention_adapter.dense_h_to_4h.bias.{self.tensor_para_rank}.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_attention_adapter.dense_4h_to_h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_attention_adapter.dense_4h_to_h.bias.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_ffn_adapter.dense_h_to_4h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_ffn_adapter.dense_h_to_4h.bias.{self.tensor_para_rank}.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_ffn_adapter.dense_4h_to_h.weight.{self.tensor_para_rank}.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]
            [
                _load_from_file(
                    f"model.layers.{layer_offset + i}.after_ffn_adapter.dense_4h_to_h.bias.bin"  # noqa: E501
                )
                for i in range(self.local_num_layers)
            ]

        assert len(self.weights) == len(
            self.expected_weight_shapes
        ), "Incorrect number of weights loaded"

class FtModuleBase:
    def __init__(self):
        self.weight = None

    @classmethod
    @abstractmethod
    def from_config(cls, config: GptInitModelParameters, **kwargs):
        raise NotImplementedError

    @abstractmethod
    def _initialize_model(self, force_init=False):
        raise NotImplementedError

    @abstractmethod
    def forward(self, *args, **kwargs):
        raise NotImplementedError

    def set_weight(self, weight: GptLayerWeights):
        old_weight_dtype = (
            self.weight.dtype if self.weight is not None else None
        )
        self.weight = weight
        if old_weight_dtype is None or old_weight_dtype != self.weight.dtype:
            self._initialize_model(force_init=True)

    @property
    def dtype(self):
        assert self.weight is not None
        return self.weight.dtype

    @property
    def device(self):
        assert self.weight is not None
        return self.weight.device

    def cuda(self, device=None):
        assert torch.cuda.is_available()
        self.weight.cuda(device)
        return self

    def to(self, device=None):
        self.weight.to(device)
        return self

    def float(self):
        self.weight.float()
        self._initialize_model(force_init=True)
        return self

    def half(self):
        self.weight.half()
        self._initialize_model(force_init=True)
        return self

    def bfloat16(self):
        self.weight.bfloat16()
        self._initialize_model(force_init=True)
        return self

class GptContextDecoder(FtModuleBase):
    def __init__(
        self,
        num_heads: int,
        size_per_head: int,
        inter_size: int,
        num_layers: int,
        tensor_para_size: int = 1,
        pipeline_para_size: int = 1,
        remove_padding: bool = True,
        shared_contexts_ratio: float = 1.0,
        layernorm_eps: float = 1e-6,
        layernorm_type: LayernormType = "pre_layernorm",
        activation_type: str = "gelu",
        has_adapters: bool = False,
        adapter_inter_size: int = 0,
        int8_mode: int = 0,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.size_per_head = size_per_head
        self.hidden_size = self.num_heads * self.size_per_head
        self.inter_size = inter_size
        self.num_layers = num_layers

        self.tensor_para_size = tensor_para_size
        self.pipeline_para_size = pipeline_para_size

        self.remove_padding = remove_padding
        self.shared_contexts_ratio = shared_contexts_ratio

        self.layernorm_eps = layernorm_eps
        self.layernorm_type = layernorm_type
        self.activation_type = activation_type
        self.has_adapters = has_adapters
        self.adapter_inter_size = adapter_inter_size

        assert int8_mode in [0, 1]
        self.int8_mode = int8_mode

        self.ft_op = None
        self.weight = None

    def __repr__(self):
        args_dict = dict(
            num_heads=self.num_heads,
            size_per_head=self.size_per_head,
            hidden_size=self.hidden_size,
            inter_size=self.inter_size,
            num_layers=self.num_layers,
            tensor_para_size=self.tensor_para_size,
            pipeline_para_size=self.pipeline_para_size,
            remove_padding=self.remove_padding,
            shared_contexts_ratio=self.shared_contexts_ratio,
            layernorm_eps=self.layernorm_eps,
            layernorm_type=self.layernorm_type,
            activation_type=self.activation_type,
            has_adapters=self.has_adapters,
            adapter_inter_size=self.adapter_inter_size,
            int8_mode=self.int8_mode,
        )
        args_str = ",\n    ".join([f"{k}: {v}" for k, v in args_dict.items()])
        return f"{self.__class__.__name__}[\n{    args_str}\n]"

    @classmethod
    def from_config(cls, config: GptInitModelParameters, **kwargs):
        return cls(
            num_heads=config.head_num,
            size_per_head=config.size_per_head,
            inter_size=4 * config.head_num * config.size_per_head,
            num_layers=config.layer_num,
            tensor_para_size=config.tensor_para_size,
            pipeline_para_size=config.pipeline_para_size,
            remove_padding=kwargs.get("remove_padding", True),
            shared_contexts_ratio=kwargs.get("shared_contexts_ratio", 1.0),
            layernorm_eps=config.layernorm_eps,
            layernorm_type=config.layernorm_type,
            activation_type=config.activation_type,
            has_adapters=config.has_adapters,
            adapter_inter_size=config.adapter_inter_size,
            int8_mode=config.int8_mode,
        )

    def _initialize_model(self, force_init=False):
        if self.weight is None:
            self.weight = GptLayerWeights(
                num_heads=self.num_heads,
                size_per_head=self.size_per_head,
                inter_size=self.inter_size,
                num_layers=self.num_layers,
                tensor_para_size=self.tensor_para_size,
                pipeline_para_size=self.pipeline_para_size,
                has_adapters=self.has_adapters,
                adapter_inter_size=self.adapter_inter_size,
                int8_mode=self.int8_mode,
            )
        if not force_init and self.ft_op is not None:
            return
        if self.ft_op is not None:
            del self.ft_op

        self.ft_op = (
            torch.classes.FasterTransformer.ParallelGptContextDecoderOp(
                self.num_heads,
                self.size_per_head,
                self.inter_size,
                self.num_layers,
                self.tensor_para_size,
                self.pipeline_para_size,
                self.layernorm_eps,
                self.layernorm_type,
                self.activation_type,
                self.has_adapters,
                self.adapter_inter_size,
                self.int8_mode,
                self.weight.weights,
                self.weight.int8_weights,
                self.weight.int8_scales,
                self.remove_padding,
            )
        )

    def forward(
        self,
        input_embeds: torch.Tensor,
        attention_mask: torch.Tensor,
        input_lengths: torch.IntTensor,
        memory_length: Optional[int] = None,
        compact_index: Optional[torch.IntTensor] = None,
        batch_to_compact_index: Optional[torch.IntTensor] = None,
        linear_bias_slopes: Optional[torch.Tensor] = None,
    ):
        """

        # Args.
            input_embeds: Tensor, (batch * beam, max_input_length, hidden_dim),
                input hidden states.
            attention_mask: Tensor, (batch * beam, max_input_length, max_input_length),
                input attention mask.
            input_lengths: (batch * beam,), input sequence lengths.
            memory_length: int, the length of memory to keep key/cache values.
            compact_index: IntTensor, (compact_batch_size,)
                The index of input sequences of a compact batch. If None, the FT op
                doesn't apply the shared context feature and as result the inference
                time may increase.
            batch_to_compact_index: IntTensor, (batch * beam,)
                The index map from the original input batch to the compact batch.
                This must be provided if compact_index is not None.
            linear_bias_slopes: (num_heads,)
                The slope per head of linear attention bias - ALiBi. If None, a base
                self attention will be performed.
        # Returns
            hidden_states: Tensor, (batch * beam, max_input_length, hidden_dim),
                decoder outputs.
            key_cache: Tensor, (num_layers, batch * beam, local_num_heads, size_per_head / x, memory_length, x), # noqa: E501
                key cache of attention of inputs.
                x = 16 / sizeof(T), memory_length = max_input_length or max_input_length + gen_length # noqa: E501
            value_cache: Tensor, (num_layers, batch * beam, local_num_heads, memory_length, hidden_dim) # noqa: E501
                value cache of attention
            last_token_hidden_states: Tensor, (batch * beam, hidden_dim)
                hidden states of the last input token.
        """
        self._initialize_model()
        # outputs: output hidden states
        (
            decoder_ouptut,
            key_cache,
            value_cache,
            last_token_hidden_states,
        ) = self.ft_op.forward(
            input_embeds,
            attention_mask,
            input_lengths,
            memory_length,
            compact_index,
            batch_to_compact_index,
            linear_bias_slopes,
        )
        return decoder_ouptut, key_cache, value_cache, last_token_hidden_states

class GptDecoder(FtModuleBase):
    def __init__(
        self,
        num_heads: int,
        size_per_head: int,
        inter_size: int,
        num_layers: int,
        tensor_para_size: int = 1,
        pipeline_para_size: int = 1,
        layernorm_eps: float = 1e-6,
        layernorm_type: LayernormType = "pre_layernorm",
        activation_type: str = "gelu",
        has_adapters: bool = False,
        adapter_inter_size: int = 0,
        int8_mode: int = 0,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.size_per_head = size_per_head
        self.hidden_size = self.num_heads * self.size_per_head
        self.inter_size = inter_size
        self.num_layers = num_layers

        self.tensor_para_size = tensor_para_size
        self.pipeline_para_size = pipeline_para_size

        self.layernorm_eps = layernorm_eps
        self.layernorm_type = layernorm_type
        self.activation_type = activation_type
        self.has_adapters = has_adapters
        self.adapter_inter_size = adapter_inter_size

        self.int8_mode = int8_mode

        self.ft_op = None
        self.weight = None

    def __repr__(self):
        args_dict = dict(
            num_heads=self.num_heads,
            size_per_head=self.size_per_head,
            hidden_size=self.hidden_size,
            inter_size=self.inter_size,
            num_layers=self.num_layers,
            tensor_para_size=self.tensor_para_size,
            pipeline_para_size=self.pipeline_para_size,
            layernorm_eps=self.layernorm_eps,
            layernorm_type=self.layernorm_type,
            activation_type=self.activation_type,
            has_adapters=self.has_adapters,
            adapter_inter_size=self.adapter_inter_size,
            int8_mode=self.int8_mode,
        )
        args_str = ",\n    ".join(
            [f"{k}: {v}" for k, v in args_dict.items()]
        )  # noqa: E501
        return f"{self.__class__.__name__}[\n    {args_str}\n]"

    @classmethod
    def from_config(cls, config: GptInitModelParameters, **kwargs):
        hidden_dim = config.head_num * config.size_per_head
        return cls(
            num_heads=config.head_num,
            size_per_head=config.size_per_head,
            inter_size=4 * hidden_dim,
            num_layers=config.layer_num,
            tensor_para_size=config.tensor_para_size,
            pipeline_para_size=config.pipeline_para_size,
            layernorm_eps=config.layernorm_eps,
            layernorm_type=config.layernorm_type,
            activation_type=config.activation_type,
            has_adapters=config.has_adapters,
            adapter_inter_size=config.adapter_inter_size,
            int8_mode=config.int8_mode,
        )

    def _initialize_model(self, force_init=False):
        if self.weight is None:
            self.weight = GptLayerWeights(
                num_heads=self.num_heads,
                size_per_head=self.size_per_head,
                inter_size=self.inter_size,
                num_layers=self.num_layers,
                tensor_para_size=self.tensor_para_size,
                pipeline_para_size=self.pipeline_para_size,
                has_adapters=self.has_adapters,
                adapter_inter_size=self.adapter_inter_size,
                int8_mode=self.int8_mode,
            )
        if not force_init and self.ft_op is not None:
            return
        if self.ft_op is not None:
            del self.ft_op
        self.ft_op = torch.classes.FasterTransformer.ParallelGptDecoderOp(
            self.num_heads,
            self.size_per_head,
            self.inter_size,
            self.num_layers,
            self.tensor_para_size,
            self.pipeline_para_size,
            self.layernorm_eps,
            self.layernorm_type,
            self.activation_type,
            self.has_adapters,
            self.adapter_inter_size,
            self.weight.int8_mode,
            self.weight.weights,
            self.weight.int8_weights,
            self.weight.int8_scales,
        )

    def forward(
        self,
        max_input_length: int,
        step: int,
        ite: int,
        input_embeds: torch.Tensor,
        sequence_lengths: torch.IntTensor,
        key_cache: torch.Tensor,
        value_cache: torch.Tensor,
        finished: torch.BoolTensor,
        total_padding_tokens: torch.IntTensor,
        masked_tokens: torch.BoolTensor,
        cache_indirection: Optional[torch.IntTensor] = None,
        linear_bias_slopes: Optional[torch.Tensor] = None,
    ):
        """

        # Args.
            max_input_length: int, maximum input context length.
            step: int, the current step index.
            ite: int, local batch iteration.
            input_embeds: Tensor, (local_batch * beam, hidden_dim),
                input hidden state to decoder.
            sequence_lengths: IntTensor, (local_batch * beam,),
                the current sequence lengths.
            key_cache: Tensor, key cache buffer.
            value_cache: Tensor, value cache buffer.
            finished: BoolTensor, (local_batch * beam,),
                whether to finish sentence generation.
            total_padding_tokens IntTensor, (local_batch * beam,),
                the number of padded tokens.
            masked_tokens: BoolTensor, (local_batch * beam, memory_length),
                a mask tensor that indicates padded tokens.
            cache_indirection: IntTensor, (local_batch * beam,),
                cache of beam positions if needed if beam > 1.
            linear_bias_slopes Tensor, (num_heads,)
                slopes head of linear position bias (ALiBi) (optional).
        # Returns
            IntTensor, (batch * beam,) output token ids.
        """

        self._initialize_model()

        outputs = self.ft_op.forward(
            max_input_length,
            step,
            ite,
            input_embeds,
            sequence_lengths,
            finished,
            total_padding_tokens,
            masked_tokens,
            key_cache,
            value_cache,
            cache_indirection,
            linear_bias_slopes,
        )
        return outputs[0]

class Gpt:
    def __init__(
        self,
        num_heads: int,
        size_per_head: int,
        num_layers: int,
        vocab_size: int,
        start_id: int,
        end_id: int,
        lib_path: PathLike,
        tensor_para_size: int = 1,
        pipeline_para_size: int = 1,
        remove_padding: bool = True,
        shared_contexts_ratio: float = 1.0,
        layernorm_eps: float = 1e-6,
        layernorm_type: LayernormType = "pre_layernorm",
        activation_type: str = "gelu",
        has_positional_encoding: bool = True,
        max_seq_len: int = 0,
        has_pre_decoder_layernorm: bool = False,
        has_post_decoder_layernorm: bool = True,
        has_adapters: bool = False,
        adapter_inter_size: int = 0,
        int8_mode: int = 0,
        inference_data_type: Optional[str] = None,
        weights_data_type: str = "fp32",
        use_fp32_to_compute_logit: bool = False,
        **kwargs,
    ):
        super().__init__()

        inference_data_type = inference_data_type or weights_data_type

        self.config = GptInitModelParameters(
            head_num=num_heads,
            size_per_head=size_per_head,
            layer_num=num_layers,
            max_seq_len=max_seq_len,
            tensor_para_size=tensor_para_size,
            vocab_size=vocab_size,
            start_id=start_id,
            end_id=end_id,
            pipeline_para_size=pipeline_para_size,
            data_type=inference_data_type,
            weights_data_type=weights_data_type,
            layernorm_eps=layernorm_eps,
            layernorm_type=layernorm_type,
            activation_type=activation_type,
            has_positional_encoding=has_positional_encoding,
            has_pre_decoder_layernorm=has_pre_decoder_layernorm,
            has_post_decoder_layernorm=has_post_decoder_layernorm,
            has_adapters=has_adapters,
            adapter_inter_size=adapter_inter_size,
            int8_mode=int8_mode,
            sparse=kwargs.get("sparse", False),
        )
        self.use_fp32_to_compute_logit = use_fp32_to_compute_logit

        self.weight = None
        self.shared_contexts_ratio = shared_contexts_ratio

        torch.classes.load_library(os.path.abspath(lib_path))

        # Embeddings to encode or decode tokens.
        hidden_dim = num_heads * size_per_head

        # Pad vocab size for FT.
        local_vocab_size = math.ceil(
            self.config.vocab_size / self.config.tensor_para_size
        )
        if self.config.data_type == "fp16":
            local_vocab_size = math.ceil(local_vocab_size / 8) * 8
        self.vocab_size_padded = (
            local_vocab_size * self.config.tensor_para_size
        )
        self.vocab_size = self.config.vocab_size

        self.decode_op = torch.classes.FasterTransformer.DynamicDecodeOp(
            self.vocab_size,
            self.vocab_size_padded,
            self.config.tensor_para_size,
            self.config.pipeline_para_size,
            torch.float,
        )

        self._parameters = {}

        def register_param(name, p):
            self._parameters[name] = p
            setattr(self, name, p)

        register_param(
            "context_decoder",
            GptContextDecoder.from_config(
                self.config,
                remove_padding=remove_padding,
                shared_contexts_ratio=shared_contexts_ratio,
                **kwargs,
            ),
        )
        register_param(
            "decoder", GptDecoder.from_config(self.config, **kwargs)
        )

        compute_dtype = to_torch_dtype(inference_data_type)

        if comm.is_pipeline_group_first():
            register_param(
                "word_embedding",
                torch.nn.Embedding(
                    self.vocab_size_padded, hidden_dim, dtype=compute_dtype
                ),
            )
            self._mask_padded_vocab_weights(self.word_embedding.weight)
            if self.config.has_positional_encoding:
                register_param(
                    "position_encoding",
                    torch.nn.Embedding(
                        self.config.max_seq_len,
                        hidden_dim,
                        dtype=compute_dtype,
                    ),
                )
            else:
                self.position_encoding = None
            if self.config.has_pre_decoder_layernorm:
                register_param(
                    "pre_decoder_layernorm",
                    torch.nn.LayerNorm(
                        hidden_dim, eps=layernorm_eps, dtype=compute_dtype
                    ),
                )
            else:
                self.pre_decoder_layernorm = None

        if comm.is_pipeline_group_last():
            if has_post_decoder_layernorm:
                register_param(
                    "post_decoder_layernorm",
                    torch.nn.LayerNorm(
                        hidden_dim, eps=layernorm_eps, dtype=compute_dtype
                    ),
                )
            else:
                self.post_decoder_layernorm = None
            self.lm_head_ctype = (
                compute_dtype
                if not self.use_fp32_to_compute_logit
                else torch.float32
            )
            register_param(
                "lm_head",
                torch.nn.Linear(
                    hidden_dim,
                    self.vocab_size_padded,
                    bias=False,
                    dtype=self.lm_head_ctype,
                ),
            )
            self._mask_padded_vocab_weights(self.lm_head.weight)

    @classmethod
    def from_config(cls, config: GptInitModelParameters, **kwargs):
        return cls(
            num_heads=config.head_num,
            size_per_head=config.size_per_head,
            num_layers=config.layer_num,
            max_seq_len=config.max_seq_len,
            tensor_para_size=config.tensor_para_size,
            vocab_size=config.vocab_size,
            start_id=config.start_id,
            end_id=config.end_id,
            pipeline_para_size=config.pipeline_para_size,
            inference_data_type=config.data_type,
            weights_data_type=config.weights_data_type,
            layernorm_eps=config.layernorm_eps,
            layernorm_type=config.layernorm_type,
            activation_type=config.activation_type,
            has_positional_encoding=config.has_positional_encoding,
            has_pre_decoder_layernorm=config.has_pre_decoder_layernorm,
            has_post_decoder_layernorm=config.has_post_decoder_layernorm,
            has_adapters=config.has_adapters,
            adapter_inter_size=config.adapter_inter_size,
            int8_mode=config.int8_mode,
            **kwargs,
        )

    def load(
        self,
        checkpoint_path: PathLike,
        inference_data_type: Optional[Union[str, torch.dtype]] = None,
        config: Optional[GptInitModelParameters] = None,
        device: Optional[Union[str, int, torch.device]] = None,
    ):

        checkpoint_path = Path(checkpoint_path)
        device = device or comm.get_device()
        config = config or self.config

        compute_dtype = to_torch_dtype(inference_data_type or self.dtype)

        self.weight = GptLayerWeights.from_config(config)
        self.weight.load(
            checkpoint_path, compute_dtype, config.weights_data_type, device
        )

        self.context_decoder.set_weight(self.weight)
        self.decoder.set_weight(self.weight)

        weight_dtype = to_numpy_dtype(config.weights_data_type)

        def _safe_load_from_bin(param: torch.nn.Parameter, fname):
            if (checkpoint_path / fname).exists():
                # np_w is 1-D array since a bin file doesn't have shape info.
                w_ = np.fromfile(checkpoint_path / fname, dtype=weight_dtype)
                param.data = (
                    torch.from_numpy(w_)
                    .reshape(param.data.shape)
                    .to(compute_dtype)
                )
            else:
                raise FileNotFoundError(f"Faile to load {fname}")

        def _safe_load_lm_head_from_bin(param, fname, ctype):
            if (checkpoint_path / fname).exists():
                shape = (
                    self.vocab_size,
                    self.config.head_num * self.config.size_per_head,
                )
                # np_w is 1-D array since a bin file doesn't have shape info.
                w_ = np.fromfile(checkpoint_path / fname, dtype=weight_dtype)
                param.data = param.data.to(ctype)
                param.data[: self.vocab_size, :] = (
                    torch.from_numpy(w_).reshape(shape).to(ctype)
                )
            else:
                print(f"Faile to load {fname}")
                torch.nn.init.normal_(param).to(compute_dtype)
            self._mask_padded_vocab_weights(param)

        # pylint:disable=line-too-long
        if comm.is_pipeline_group_first():
            _safe_load_lm_head_from_bin(
                self.word_embedding.weight, "model.wte.bin", compute_dtype
            )
            self._mask_padded_vocab_weights(self.word_embedding.weight)
            if self.position_encoding is not None:
                _safe_load_from_bin(
                    self.position_encoding.weight, "model.wpe.bin"
                )
            if self.pre_decoder_layernorm is not None:
                _safe_load_from_bin(
                    self.pre_decoder_layernorm.weight,
                    "model.pre_decoder_layernorm.weight.bin",
                )
                _safe_load_from_bin(
                    self.pre_decoder_layernorm.bias,
                    "model.pre_decoder_layernorm.bias.bin",
                )
        if comm.is_pipeline_group_last():
            if self.post_decoder_layernorm is not None:
                _safe_load_from_bin(
                    self.post_decoder_layernorm.weight,
                    "model.final_layernorm.weight.bin",
                )
                _safe_load_from_bin(
                    self.post_decoder_layernorm.bias,
                    "model.final_layernorm.bias.bin",
                )
            if (checkpoint_path / "model.lm_head.weight.bin").exists():
                _safe_load_lm_head_from_bin(
                    self.lm_head.weight,
                    "model.lm_head.weight.bin",
                    self.lm_head_ctype,
                )
            else:
                if self.use_fp32_to_compute_logit:
                    _safe_load_lm_head_from_bin(
                        self.lm_head.weight, "model.wte.bin", torch.float32
                    )
                else:
                    # In this branch we can share the pre and post
                    # decoder embeddings, but ONLY pipeline size is 1.
                    # When pipeline size > 1, these two weights will end up on
                    # different GPUs, so we must load the
                    # post decoder weight again (else case).
                    if comm.get_pipeline_para_size() == 1:
                        self.lm_head.weight = self.word_embedding.weight
                    else:
                        _safe_load_lm_head_from_bin(
                            self.lm_head.weight, "model.wte.bin", compute_dtype
                        )

        self.to(device)

    @property
    def dtype(self):
        assert self.weight is not None
        return self.weight.dtype

    @property
    def device(self):
        assert self.weight is not None
        return self.weight.device

    def cuda(self, device=None):
        assert torch.cuda.is_available()
        for name, param in self._parameters.items():
            setattr(self, name, param.cuda(device))
        return self

    def to(self, device=None):
        for name, param in self._parameters.items():
            setattr(self, name, param.to(device))
        return self

    def float(self):
        for name, param in self._parameters.items():
            setattr(self, name, param.float())
        return self

    def half(self):
        for name, param in self._parameters.items():
            setattr(self, name, param.half())
        return self

    def bfloat16(self):
        for name, param in self._parameters.items():
            setattr(self, name, param.bfloat16())
        return self

    def _mask_padded_vocab_weights(self, weight: torch.Tensor):
        assert self.vocab_size_padded >= self.vocab_size
        if self.vocab_size_padded > self.vocab_size:
            weight.data[self.vocab_size :, ...] = 0  # noqa: E203

    def generate_pad_mask(self, input_lengths, memory_length, init_step=0):
        """Generate a pad mask tensor.

        # Args.
            input_lengths: (batch_size * beam_width,), input lengths
            memory_length: the length of key/value cache memory.
            init_step: int, initial step.
        # Return
            masked_tokens: BoolTensor,
                (batch_size * beam_width, memory_length),
                True if init_step + input_length[i] <= j <
                    init_step + max_input_length,
                where i is a batch-beam index and j is a time step
                modulo by memory_length.
        """
        max_input_length = input_lengths.max()
        input_lengths = input_lengths.unsqueeze(1)
        shift = init_step % memory_length
        step_indices = torch.arange(
            init_step, init_step + memory_length, device=input_lengths.device
        )
        step_indices = (
            step_indices.roll(shift)
            .unsqueeze(0)
            .tile(input_lengths.shape[0], 1)
        )
        masked_tokens = torch.logical_and(
            step_indices >= input_lengths,
            step_indices < init_step + max_input_length,
        )
        return masked_tokens

    def get_local_batch_size(self, batch_size):
        """Get a local batch size by the same way that FT Gpt does."""
        local_batch_size = batch_size
        pp_size = self.decoder.pipeline_para_size
        if pp_size > 1:
            if local_batch_size % pp_size == 0:
                local_batch_size //= pp_size
            while local_batch_size > 1024 and local_batch_size % 2 == 0:
                local_batch_size //= 2
        return local_batch_size

    @torch.no_grad()
    def generate(
        self,
        input_token_ids: torch.IntTensor,
        input_lengths: torch.IntTensor,
        gen_length: int,
        eos_token_id: Optional[int] = None,
        local_batch_size: Optional[int] = None,
        beam_width: int = 1,
        top_k: Optional[torch.IntTensor] = None,
        top_p: Optional[torch.FloatTensor] = None,
        top_p_decay: Optional[torch.FloatTensor] = None,
        top_p_min: Optional[torch.FloatTensor] = None,
        top_p_reset_ids: Optional[torch.IntTensor] = None,
        temperature: Optional[torch.FloatTensor] = None,
        repetition_penalty: Optional[torch.FloatTensor] = None,
        presence_penalty: Optional[torch.FloatTensor] = None,
        min_length: Optional[torch.IntTensor] = None,
        len_penalty: Optional[torch.FloatTensor] = None,
        beam_search_diversity_rate: Optional[torch.FloatTensor] = None,
        stop_words_list: Optional[torch.IntTensor] = None,
        bad_words_list: Optional[torch.IntTensor] = None,
        sequence_limit_lengths: Optional[torch.IntTensor] = None,
        random_seed: Optional[torch.LongTensor] = None,
        memory_length: Optional[int] = None,
        return_output_length: bool = False,
        return_log_probs: bool = False,
    ):
        """

        # Args.
            input_token_ids: IntTensor, (batch_size, max_input_length),
                input hidden state to decoder.
            input_lengths: IntTensor, (batch_size),
                the lengths of input context sequences.
            gen_length: int, the number of tokens to generate.
            local_batch_size: int, optional, a batch size of
                local iteration. (disabled)
            eos_token_id: int, eos token id.
            beam_width: int, number of beams for beam search.
                If 1, sampling decode will be used.
            top_k: IntTensor, (batch_size,) top-k sampling.
                The number of most probable tokens to keep
                for sampling per sentence in a batcch.
            top_p: FloatTensor, (batch_size,), top-p sampling.
                The cumulative probability
                of to filter the set of most probable tokens.
            top_p_decay: FloatTensor, (batch_size,)
                The decay of top-p value for top_p sampling.
            top_p_min: FloatTensor, (batch_size,)
                The minimum top p values in top-p decaying.
            top_p_reset_ids: IntTensor, (batch_size,)
                reset ids for resetting top_p values for top p sampling
            temperature: FloatTensor, (batch_size,),
                The temperature value for smoothing the logit distribution.
            repetition_penalty: FloatTensor, (batch_size,),
                The repetition penalty.
            presence_penalty: FloatTensor, (batch_size,),
                The presence penalty, which is exclusive with
                repetition_penalty.
                Only one of repetition and presence penalties is allowed.
            min_length: IntTensor, (batch_size,),
                Minimum length for each sentences. EOS is masked if length is
                below min.
            len_penalty: FloatTensor, (batch_size,)
                The exponent of the length penalty of beam scores.
            beam_search_diversity_rate: FloatTensor, (batch_size,),
                The diversity rate of beam search.
            stop_words_list: IntTensor, (batch_size, 2, stop_words_length)
                When FT generates words in this list, it will stop the
                generation. An extension of stop id.
            bad_words_list IntTensor, (batch_size, 2, bad_words_length)
                The words in the list will never be sampled.
            sequence_limit_lengths: IntTensor, (batch_size,), The maximum
                length of a generated sequence.
            memory_length: int, the length of cache memory. If None, it will
                be max_input_length + gen_length.
        # Returns
            IntTensor, (batch_size, beam_width, max_seq_length) output
            token ids.
        """
        assert (
            self.weight is not None
        ), "Please call load() first to initialize weights."

        input_token_ids = input_token_ids.type(torch.int32).to(self.device)
        input_lengths = input_lengths.type(torch.int32).to(self.device)

        batch_size = len(input_token_ids)
        max_input_length = input_token_ids.shape[-1]
        max_seq_length = max_input_length + gen_length
        memory_length = memory_length or max_seq_length

        # TODO: Enable local batch later. We currently disable local batching due to # noqa: E501
        #   an input mismatch issue of FT's decode_op: FT's decode_op requires logits # noqa: E501
        #   of shape (batch_size, ...) but we have logits of shape (local_batch_size, ...) # noqa: E501
        #   After fixing FT's side, we will enable local batch.
        # local_batch_size = local_batch_size or self.get_local_batch_size(batch_size) # noqa: E501
        # num_local_batches, last_chunk = divmod(batch_size, local_batch_size)
        # if last_chunk > 0:
        #     num_local_batches += 1
        assert local_batch_size is None or local_batch_size == batch_size
        local_batch_size = batch_size
        num_local_batches = 1

        device = self.device

        eos_token_id = (
            eos_token_id if eos_token_id is not None else self.config.end_id
        )
        assert (
            eos_token_id is not None
        ), "eos_token-id must be specified in generation."
        eos_token_ids = eos_token_id * torch.ones(
            batch_size, dtype=torch.int32, device=device
        )
        assert repetition_penalty is None or presence_penalty is None, (
            "Found ambiguous parameters repetition_penalty and "
            "presence_penalty which are mutually exclusive. "
            "Please provide one of repetition_penalty and presence_penalty."
        )
        # Setup decoder_op prior to calling the forward function.
        self.decode_op.setup(
            batch_size,
            beam_width,
            top_k,
            top_p,
            temperature,
            repetition_penalty,
            presence_penalty,
            min_length,
            len_penalty,
            beam_search_diversity_rate,
            random_seed,
            top_p_decay,
            top_p_min,
            top_p_reset_ids,
        )

        # Prepare input and output arguments.
        if beam_width > 1:
            # Tiling for beam search.
            input_token_ids = input_token_ids.repeat(1, beam_width).view(
                batch_size * beam_width, -1
            )
            input_lengths = (
                input_lengths.view(-1, 1).repeat(1, beam_width).view(-1)
            )
            if sequence_limit_lengths is not None:
                sequence_limit_lengths = (
                    sequence_limit_lengths.view(-1, 1)
                    .repeat(1, beam_width)
                    .view(-1)
                )
            # src/tgt cache indirections.
            cache_indirection = torch.zeros(
                (2, batch_size, beam_width, memory_length),
                dtype=torch.int32,
                device=device,
            )
            parent_ids = torch.zeros(
                max_seq_length,
                batch_size * beam_width,
                dtype=torch.int32,
                device=device,
            )
        else:
            cache_indirection = None
            src_cache_indirection = None
            tgt_cache_indirection = None
            parent_ids = None

        pad_lengths = max_input_length - input_lengths
        # Since tril() doesn't support bf16 dtype,
        # we create of bool type and then cast it to dtype.
        attention_mask = (
            torch.ones(
                (max_input_length, max_input_length),
                dtype=torch.bool,
                device=device,
            )
            .tril()
            .unsqueeze(0)
            .tile(input_token_ids.shape[0], 1, 1)
            .to(self.dtype)
        )
        for b, input_length in enumerate(input_lengths):
            attention_mask[b, input_length:, ...] = 0
        masked_tokens = self.generate_pad_mask(input_lengths, memory_length)
        finished = torch.zeros_like(input_lengths).bool()
        sequence_lengths = (max_input_length - 1) * torch.ones_like(
            input_lengths
        )

        if return_log_probs or beam_width > 1:
            cum_log_probs = torch.zeros(batch_size * beam_width, device=device)
            output_log_probs = torch.zeros(
                (gen_length, batch_size * beam_width), device=device
            )
        else:
            cum_log_probs = None
            output_log_probs = None

        # Contiguous buffer for each decode_op step,
        # it will be transposed tensor for the final output.
        output_token_ids = torch.zeros(
            (max_seq_length, batch_size * beam_width),
            dtype=torch.int32,
            device=device,
        )
        output_token_ids[:max_input_length, ...] = input_token_ids.T

        if comm.is_pipeline_group_first():
            # Prepare input tensors of decoder.
            input_embeds = self.word_embedding(input_token_ids)
            if self.position_encoding is not None:
                position_ids = torch.arange(
                    0, max_input_length, dtype=torch.int, device=device
                )
                position_ids = position_ids.unsqueeze(0).view(
                    -1, max_input_length
                )
                input_embeds += self.position_encoding(position_ids)
            if self.pre_decoder_layernorm is not None:
                input_embeds = self.pre_decoder_layernorm(input_embeds)
        else:
            # Dummy input_embeds
            input_embeds = torch.empty(
                size=(
                    batch_size * beam_width,
                    max_input_length,
                    self.context_decoder.hidden_size,
                ),
                dtype=self.context_decoder.dtype,
                device=device,
            )

        use_shared_contexts = (
            (self.shared_contexts_ratio > 0.0)
            and (max_input_length >= 1)
            and (batch_size > 1)
        )
        batch_to_compact, compact_to_batch = None, None
        if use_shared_contexts:
            find_context_duplications = (
                torch.ops.fastertransformer.find_context_duplications
            )
            batch_to_compact, compact_to_batch = find_context_duplications(
                input_token_ids
            )
            use_shared_contexts = (
                compact_to_batch.shape[0]
                <= self.shared_contexts_ratio * batch_size
            )

            if not use_shared_contexts:
                batch_to_compact, compact_to_batch = None, None

        profiler.start("ft-context-decoder")
        (
            _,
            k_cache,
            v_cache,
            last_token_hidden_states,
        ) = self.context_decoder.forward(
            input_embeds=input_embeds,
            attention_mask=attention_mask,
            input_lengths=input_lengths,
            memory_length=memory_length,
            batch_to_compact_index=batch_to_compact,
            compact_index=compact_to_batch,
        )
        profiler.stop("ft-context-decoder")

        for step in range(max_input_length, max_seq_length):
            src_indir_idx = (step - max_input_length) % 2
            tgt_indir_idx = 1 - src_indir_idx

            is_generation_done = torch.tensor(
                [True], dtype=torch.bool, device=device
            )
            for ite in range(num_local_batches):
                # The indices of the current local batch-beam.
                bbidx = range(
                    ite * local_batch_size * beam_width,
                    min(
                        (ite + 1) * local_batch_size * beam_width,
                        batch_size * beam_width,
                    ),
                )
                if cache_indirection is not None:
                    bidx = range(
                        ite * local_batch_size,
                        min((ite + 1) * local_batch_size, batch_size),
                    )
                    src_cache_indirection = cache_indirection[
                        src_indir_idx, bidx, ...
                    ]
                    tgt_cache_indirection = cache_indirection[
                        tgt_indir_idx, bidx, ...
                    ]

                if step == max_input_length:
                    hidden_states = last_token_hidden_states[bbidx, ...]
                else:
                    if comm.is_pipeline_group_first():
                        input_embeds = self.word_embedding(
                            output_token_ids[step - 1, bbidx]
                        )
                        if self.position_encoding is not None:
                            position_ids = (step - 1) * torch.ones_like(
                                pad_lengths[bbidx]
                            )
                            input_embeds += self.position_encoding(
                                position_ids
                            )
                        if self.pre_decoder_layernorm is not None:
                            input_embeds = self.pre_decoder_layernorm(
                                input_embeds
                            )
                    else:
                        # Dummy input_imbeds
                        input_embeds = torch.empty(
                            size=(len(bbidx), self.decoder.hidden_size),
                            dtype=self.decoder.dtype,
                            device=device,
                        )

                    profiler.start("ft-decoder")
                    hidden_states = self.decoder.forward(
                        max_input_length=max_input_length,
                        step=step,
                        ite=ite,
                        input_embeds=input_embeds,
                        sequence_lengths=sequence_lengths[bbidx],
                        key_cache=k_cache,
                        value_cache=v_cache,
                        finished=finished[bbidx],
                        total_padding_tokens=pad_lengths[bbidx],
                        cache_indirection=src_cache_indirection,
                        masked_tokens=masked_tokens[bbidx, ...],
                    )
                    profiler.stop("ft-decoder")

                if comm.is_pipeline_group_last():
                    if self.post_decoder_layernorm is not None:
                        hidden_states = self.post_decoder_layernorm(
                            hidden_states
                        )

                    # We use logits of fp32 type to avoid overflow issue.
                    if self.use_fp32_to_compute_logit:
                        # The FT GPT op internally uses FP32 compute type
                        # for matrix multiplication.
                        # This will produce the same result with the
                        # end-to-end FT's GPT op.
                        logits = torch.nn.functional.linear(
                            hidden_states.float(), self.lm_head.weight
                        )
                    else:
                        logits = self.lm_head(hidden_states).float()

                    profiler.start("ft-decode")
                    should_stop = self.decode_op.forward(
                        logits.view(batch_size, beam_width, -1),
                        step,
                        max_input_length,
                        ite,
                        local_batch_size,
                        eos_token_ids,
                        top_k,
                        top_p,
                        temperature,
                        repetition_penalty,
                        presence_penalty,
                        min_length,
                        len_penalty,
                        beam_search_diversity_rate,
                        top_p_decay,
                        top_p_min,
                        top_p_reset_ids,
                        None,
                        input_lengths,
                        sequence_limit_lengths,
                        stop_words_list,
                        bad_words_list,
                        src_cache_indirection,
                        output_token_ids.view(-1, batch_size, beam_width),
                        finished,
                        sequence_lengths,
                        cum_log_probs,
                        output_log_probs,
                        parent_ids,
                        tgt_cache_indirection,
                    )
                    profiler.stop("ft-decode")
                    is_generation_done &= should_stop

            # Broadcast from the last pipeline node if needed.
            profiler.start("ft-bcast")
            tensors_to_bcast = [
                output_token_ids[step, ...],
                finished,
                sequence_lengths,
                is_generation_done,
            ]
            if beam_width > 1:
                tensors_to_bcast.append(tgt_cache_indirection)
            self.decode_op.broadcast_from_last_pipeline(tensors_to_bcast)
            profiler.stop("ft-bcast")

            if is_generation_done or finished.all():
                break

        # Transpose (L, batch, beam) -> (batch, beam, L)
        output_token_ids = output_token_ids.view(
            -1, batch_size, beam_width
        ).permute(1, 2, 0)

        # Increase sequence_length by 1 because the sequence length of time step t is t - 1. # noqa: E501
        sequence_lengths += 1

        # Outputs
        output_dict = dict(output_token_ids=output_token_ids)
        if return_output_length:
            output_dict["output_lengths"] = sequence_lengths
        if return_log_probs:
            output_dict["cum_log_probs"] = cum_log_probs
            output_dict["output_log_probs"] = output_log_probs
        return output_dict


"""Given two strings, write a function to return True if the strings are anagrams of each other and False if they are not.
A word is not an anagram of itself.
"""

def is_anagram(string1, string2):
    if string1 == string2 or len(string1) != len(string2):
        return False
    string1_list = sorted(string1)
    string2_list = sorted(string2)
    return string1_list == string2_list

string_1 = "listen"
string_2 = "silent"
print(is_anagram(string_1, string_2))

string_1 = "banana"
string_2 = "bandana"
print(is_anagram(string_1, string_2))

string_1 = "banana"
string_2 = "banana"
print(is_anagram(string_1, string_2))