from flask import Blueprint, g, request, render_template, session, flash, \
    redirect, url_for, abort, current_app
from flask_openid import OpenID
from flask_oauthlib.client import OAuth
from sqlalchemy.exc import IntegrityError

from sphinxweb import support
from sphinxweb.models import User, db_session

auth = Blueprint(
    'auth', __name__,
    template_folder='templates',
    static_folder='static',
    static_url_path='/auth/static')

# setup flask-openid
oid = OpenID()

oid_providers = [
    ('Launchpad', 'launchpad.png', 'https://login.launchpad.net/'),
]

oauth = OAuth()
google = oauth.remote_app(
    'google',
    request_token_params={'scope': 'email'},
    base_url='https://www.googleapis.com/oauth2/v1/',
    request_token_url=None,
    access_token_method='POST',
    access_token_url='https://accounts.google.com/o/oauth2/token',
    authorize_url='https://accounts.google.com/o/oauth2/auth',
    app_key='GOOGLE',
)

github = oauth.remote_app(
    'github',
    request_token_params={'scope': 'user:email'},
    base_url='https://api.github.com/',
    request_token_url=None,
    access_token_method='POST',
    access_token_url='https://github.com/login/oauth/access_token',
    authorize_url='https://github.com/login/oauth/authorize',
    app_key='GITHUB',
)

oauth2_providers = {
    'google': {'api': google, 'resp': 'userinfo', 'icon': 'google.png'},
    'github': {'api': github, 'resp': 'user', 'icon': 'github.png'},
}

@auth.record_once
def on_register(state):
    oauth.init_app(state.app)

@auth.route('/_login', methods=['GET', 'POST'])
@oid.loginhandler
def login():
    """Does the login via OpenID. Has to call into `oid.try_login`
    to start the OpenID machinery.
    """
    # if we are already logged in, go back to were we came from
    if g.user is not None:
        return redirect(oid.get_next_url())
    if request.method == 'POST':
        openid = request.form.get('openid')
        if openid:
            return oid.try_login(openid, ask_for=['email', 'fullname',
                                                  'nickname'])
    return render_template('login.html', next=oid.get_next_url(),
                           error=oid.fetch_error(),
                           oid_providers=oid_providers,
                           oauth2_provider_icons=[(k, v['icon']) for k, v in oauth2_providers.items()],
                           )

@oid.after_login
def create_or_login(resp):
    """This is called when login with OpenID succeeded and it's not
    necessary to figure out if this is the users's first login or not.
    This function has to redirect otherwise the user will be presented
    with a terrible URL which we certainly don't want.
    """
    session['token_key'] = token_key = 'openid'
    session[token_key] = resp.identity_url
    user = User.query.filter_by(openid=resp.identity_url).first()
    if user is not None:
        flash(u'Successfully logged in.')
        g.user = user
        return redirect(oid.get_next_url())
    return redirect(url_for('.create_profile', next=oid.get_next_url(),
                            name=resp.fullname or resp.nickname,
                            email=resp.email))

@auth.route('/_login_google')
def login_oauth2():
    key = request.args.get('key')
    if key not in oauth2_providers:
        return redirect(url_for('.login'))

    provider = oauth2_providers[key]['api']
    session['token_key'] = key
    return provider.authorize(callback=url_for('.oauth2callback', _external=True))

@auth.route('/_oauth2callback')
def oauth2callback():
    key = session['token_key']
    provider = oauth2_providers[key]['api']
    resp_field = oauth2_providers[key]['resp']

    resp = provider.authorized_response()
    if resp is None:
        return 'Access denied: reason=%s error=%s' % (
            request.args['error_reason'],
            request.args['error_description']
        )
    session[key] = resp['access_token']
    me = provider.get(resp_field)
    session['user_id'] = '{key}/{id}'.format(key=key, id=me.data['id'])
    user = User.query.filter_by(openid=session['user_id']).first()

    if user is not None:
        flash(u'Successfully logged in.')
        g.user = user
        return redirect(url_for('docs.index'))
    return redirect(url_for('.create_profile', next=url_for('docs.index'),
                            name=me.data['name'], email=me.data['email']))

@google.tokengetter
def get_google_oauth_token():
    return (session.get('google'), current_app.secret_key)

@github.tokengetter
def get_github_oauth_token():
    return (session.get('github'), current_app.secret_key)

@auth.route('/_create_profile', methods=['GET', 'POST'])
def create_profile():
    """If this is the user's first login, the create_or_login function
    will redirect here so that the user can set up his profile.
    """
    if g.user is not None or 'token_key' not in session or session['token_key'] not in session:
        return redirect(url_for('docs.index'))
    if request.method == 'POST':
        name = request.form['name']
        email = request.form['email']
        if not name:
            flash(u'Error: you have to provide a name.')
        elif '@' not in email:
            flash(u'Error: you have to enter a valid email address.')
        else:
            db_session.query(User).filter(User.name == name)
            try:
                db_session.add(User(name, email, session['user_id']))
                db_session.commit()
            except IntegrityError:
                flash(u'Error: user name or email address already taken.')
            else:
                flash(u'Profile successfully created.')
                return redirect(oid.get_next_url())
    return render_template('create_profile.html', next_url=oid.get_next_url())

@auth.route('/_profile', methods=['GET', 'POST'])
def edit_profile():
    """Update profile info."""
    if g.user is None:
        abort(401)
    form = dict(name=g.user.name, email=g.user.email)
    if request.method == 'POST':
        if 'delete' in request.form:
            db_session.delete(g.user)
            db_session.commit()
            session[session['token_key']] = None
            flash(u'Profile successfully deleted.')
            return redirect(url_for('docs.index'))
        form['name'] = request.form['name']
        form['email'] = request.form['email']
        if not form['name']:
            flash(u'Error: you have to provide a name.')
        elif '@' not in form['email']:
            flash(u'Error: you have to enter a valid email address.')
        else:
            support.update_username(g.user.name, form['name'])
            try:
                g.user.name = form['name']
                g.user.email = form['email']
                db_session.commit()
            except IntegrityError:
                flash(u'Error: user name or email address already taken.')
            else:
                flash(u'Profile successfully updated.')
                return redirect(url_for('.edit_profile'))
    return render_template('edit_profile.html', form=form)

@auth.route('/_logout')
def logout():
    session.pop('user_id', None)
    session.pop(session.pop('token_key', None), None)
    flash(u'You were logged out.')
    return redirect(oid.get_next_url())


import os
import json
import re
from transformers import PreTrainedTokenizerFast
import argparse

PATH_PARENT = os.path.abspath(os.path.join(os.getcwd(), os.pardir))
DATA_DIR = os.path.join(PATH_PARENT, "ar5iv")
QA_DIR = os.path.join(PATH_PARENT, "ArXivQA", "papers")

CLEAN = "clean_dset.json"  # dedup id --> actual id
NUM_TOKENS = "paper_id_num_token.json"  # actual id --> num tokens

parser = argparse.ArgumentParser()
parser.add_argument(
    "--latex", action="store_true", help="input is latex (default is markdown)"
)
args = parser.parse_args()

# QA markdown to json
def parse_qa(filename, tokenizer):
    paper = {}
    qas = []
    num_qa_tok = 0
    with open(filename, "r") as f_qa:
        qa = {}
        buffer = []
        for line in f_qa.readlines():
            if re.search(r"^#\s", line):
                sep_index = line.index("](https://")
                paper["title"] = line[3:sep_index].strip()
                paper["link"] = line[sep_index + 2 : -2].strip()
            elif re.search(r"^##\s", line):
                if "question" in qa:
                    qa["answer"] = "".join(buffer).strip()
                    num_qa_tok += len(tokenizer.tokenize(qa["answer"]))
                    qas.append(qa)
                    qa = {}
                qa["question"] = line[2:].strip()
                num_qa_tok += len(tokenizer.tokenize(qa["question"]))
                buffer = []
            else:
                buffer.append(line)
    paper["Q&A"] = qas

    return paper, num_qa_tok

def parse_paper(filename, tokenizer):
    paper = {}
    num_article_tok = 0

    with open(filename, "r") as f:
        article = f.read()
        paper["article"] = article
        num_article_tok += len(tokenizer.tokenize(article))

    return paper, num_article_tok

if __name__ == "__main__":
    ext = "tex" if args.latex else "md"
    clean_link = {}
    with open(CLEAN, "r") as f:
        clean_link = json.load(f)

    tokenizer = PreTrainedTokenizerFast(tokenizer_file="phi-1_5.json")
    total_num_qa_tok = 0
    total_article_num_tok = 0
    for id, actual in clean_link.items():
        print(id)
        ym, num = actual.split(".")
        paper_path = os.path.join(DATA_DIR, ym, num)

        qa_file = os.path.join(QA_DIR, f"{actual}.md")
        paper_qas, qa_num_tok = parse_qa(qa_file, tokenizer)

        article_file = os.path.join(DATA_DIR, ym, num, f"{actual}.{ext}")
        article, article_num_tok = parse_paper(article_file, tokenizer)

        total_num_qa_tok += qa_num_tok
        total_article_num_tok += article_num_tok

        paper_qas.update(article)

        with open(
            os.path.join(
                paper_path,
                f"{actual}_tex.json" if args.latex else f"{actual}.json",
            ),
            "w",
        ) as f_qa:
            json.dump(paper_qas, f_qa, indent=4)

    print(f"total {len(clean_link)} papers")
    print(f"total {total_num_qa_tok} tokens for Q&A")
    print(
        f"total {total_article_num_tok + total_num_qa_tok} tokens for paper + Q&A"
    )


import math

from absl.testing import absltest
from absl.testing import parameterized

from dp_accounting import dp_event
from dp_accounting import privacy_accountant_test
from dp_accounting.pld import pld_privacy_accountant

class PldPrivacyAccountantTest(privacy_accountant_test.PrivacyAccountantTest,
                               parameterized.TestCase):

  def _make_test_accountants(self):
    return [pld_privacy_accountant.PLDAccountant()]

  @parameterized.parameters(
      dp_event.GaussianDpEvent(1.0),
      dp_event.SelfComposedDpEvent(dp_event.GaussianDpEvent(1.0), 6),
      dp_event.ComposedDpEvent(
          [dp_event.GaussianDpEvent(1.0),
           dp_event.GaussianDpEvent(2.0)]),
      dp_event.PoissonSampledDpEvent(0.1, dp_event.GaussianDpEvent(1.0)),
      dp_event.ComposedDpEvent([
          dp_event.PoissonSampledDpEvent(0.1, dp_event.GaussianDpEvent(1.0)),
          dp_event.GaussianDpEvent(2.0)
      ]))
  def test_supports_gaussian(self, event):
    pld_accountant = pld_privacy_accountant.PLDAccountant()
    self.assertTrue(pld_accountant.supports(event))

  @parameterized.parameters(0, -1)
  def test_non_positive_composition_value_error(self, count):
    event = dp_event.GaussianDpEvent(1.0)
    accountant = pld_privacy_accountant.PLDAccountant()
    with self.assertRaises(ValueError):
      accountant.compose(event, count)

  def test_gaussian_basic(self):
    gaussian_event = dp_event.GaussianDpEvent(noise_multiplier=math.sqrt(3))
    accountant = pld_privacy_accountant.PLDAccountant()
    accountant.compose(gaussian_event, 1)
    accountant.compose(gaussian_event, 2)

    exact_epsilon = 1
    exact_delta = 0.126936
    self.assertAlmostEqual(
        accountant.get_delta(exact_epsilon), exact_delta, delta=1e-3)
    self.assertAlmostEqual(
        accountant.get_epsilon(exact_delta), exact_epsilon, delta=1e-3)

  def test_poisson_subsampled_gaussian(self):
    subsampled_gaussian_event = dp_event.PoissonSampledDpEvent(
        0.2, dp_event.GaussianDpEvent(noise_multiplier=0.5))
    accountant = pld_privacy_accountant.PLDAccountant()
    accountant.compose(subsampled_gaussian_event, 1)
    accountant.compose(subsampled_gaussian_event, 2)

    exact_epsilon = 1
    expected_delta = 0.15594
    self.assertAlmostEqual(
        accountant.get_delta(exact_epsilon), expected_delta, delta=1e-3)
    self.assertAlmostEqual(
        accountant.get_epsilon(expected_delta), exact_epsilon, delta=1e-3)

  def test_self_composed_subsampled_gaussian(self):
    event = dp_event.SelfComposedDpEvent(
        dp_event.PoissonSampledDpEvent(0.2, dp_event.GaussianDpEvent(0.5)), 3)
    accountant = pld_privacy_accountant.PLDAccountant()
    accountant.compose(event)

    exact_epsilon = 1
    expected_delta = 0.15594
    self.assertAlmostEqual(
        accountant.get_delta(exact_epsilon), expected_delta, delta=1e-3)
    self.assertAlmostEqual(
        accountant.get_epsilon(expected_delta), exact_epsilon, delta=1e-3)

if __name__ == '__main__':
  absltest.main()


import base64
import datetime
import hashlib
import hmac

GMT_FORMAT = '%a, %d %b %Y %H:%M:%S GMT'

class Auth:
    def __init__(self, secretId: str, secretKey: str, headers: dict = None):
        self.secretKey = secretKey
        self.secretId = secretId
        self.headers = headers if headers else {}

    def generateHeaders(self) -> dict:
        self.headers['Date'] = datetime.datetime.utcnow().strftime(GMT_FORMAT)
        authStr = f'hmac id="{self.secretId}", algorithm="hmac-sha1", headers="{" ".join(key.lower() for key in self.headers.keys())}", signature='
        signStr = '\n'.join(
            f'{key.lower()}: {value}' for key, value in self.headers.items())
        originSign = hmac.new(self.secretKey.encode(), signStr.encode(),
                              hashlib.sha1).digest()
        b64Sign = base64.b64encode(originSign).decode()
        authorization = authStr + f'"{b64Sign}"'
        self.headers['Authorization'] = authorization
        return self.headers


#!/usr/bin/env python
# coding: utf-8

# # Urban Sound Classification

# In[89]:

# In[95]:

import IPython.display as ipd
ipd.Audio('D:\\Urban Sound Classification\\train\\Train\\2022.wav')

# In[90]:

import numpy as np
import pandas as pd
import librosa
import librosa.display
import glob
import matplotlib.pyplot as plt
import seaborn as sns
from tqdm import tqdm
import os

# In[2]:

test = pd.read_csv("D:\\Urban Sound Classification\\test\\test.csv", index_col = 0)
train = pd.read_csv("D:\\Urban Sound Classification\\train\\train.csv", index_col = 0)
exclude = pd.read_csv("C:\\Users\\Zara\\Desktop\\exclude.csv", index_col = 0)
exclude_test = pd.read_csv("C:\\Users\\Zara\\Desktop\\exclude_test.csv", index_col = 0)

# In[3]:

#masking mismatched or missing indexes from the train and test csv files

a_index = train.index
b_index = exclude.index
mask = ~a_index.isin(b_index)
train = train.loc[mask]
train = pd.DataFrame(train)
train = train.reset_index()
train = train.reset_index(drop = True)
train = pd.DataFrame(train)
train.head(7)

# In[4]:

a_index = test.index
b_index = exclude_test.index
mask = ~a_index.isin(b_index)
test = test.loc[mask]
test = pd.DataFrame(test)
test = test.reset_index()
test = test.reset_index(drop = True)
test = pd.DataFrame(test)
test.head(7)

# In[93]:

y, sr = librosa.load('D:\\Urban Sound Classification\\train\\Train\\2022.wav')

# In[94]:

plt.figure(figsize=(12, 4))
plt.title('Waveplot: Children Playing')
librosa.display.waveplot(y, sr=sr)

# In[75]:

#counts of items in each category

train['Class'].value_counts()
plt.bar(train['Class'].value_counts().index, train['Class'].value_counts(), align='center', data=train)
plt.title('Class Distribution')
plt.xticks(rotation='vertical')
plt.xlabel('Class')
plt.ylabel('Frequency')
print('Class Distribution:\n', train['Class'].value_counts())

# In[82]:

#MFCC of children playing
mfccs = librosa.feature.mfcc(y, sr, n_mfcc=40)
plt.figure(figsize=(10,4))
librosa.display.specshow(mfccs, x_axis='time')
plt.colorbar()
plt.title('MFCC')
plt.tight_layout()

# In[81]:

#melspectrogram of children playing
melspectrogram =librosa.feature.melspectrogram(y=y, sr=sr, n_mels=40,fmax=8000)
plt.figure(figsize=(10,4))
librosa.display.specshow(librosa.power_to_db(melspectrogram,ref=np.max),y_axis='mel', fmax=8000,x_axis='time')
plt.colorbar(format='%+2.0f dB')
plt.title('Mel spectrogram')
plt.tight_layout()

#For explanation purpose.
#this function outputs the indexes of the files that cannot be read or are corrupt
#the output of the function was used to create the exclusion and exclusion_test files

def parser(row):
   # function to load files and extract features
    file_name = os.path.join(os.path.abspath(data_dir),str(row.ID)+'.wav')

   # handle exception to check if there isn't a file which is corrupted
    try:
      # here kaiser_fast is a technique used for faster extraction
        X, sample_rate = librosa.load(file_name, res_type='kaiser_fast')
      # we extract mfcc feature from data
        mfccs = np.mean(librosa.feature.mfcc(y=X, sr=sample_rate, n_mfcc=40).T,axis=0)
    except Exception as e:
        print("Error encountered while parsing file: ", file)
        return None, None

err = train.apply(parser, axis=1)
err.columns = ['feature', 'label']
# In[5]:

#extracting features from train set
#taking a mean of the mfccs

features=[]
labels=[]
for i in range(len(train)):
    data_dir = "D:\\Urban Sound Classification\\train\\Train\\"
    filename = os.path.join(os.path.abspath(data_dir)+'\\'+str(train.ID[i])+'.wav')
    x, sample_rate = librosa.load(filename, res_type='kaiser_fast')
    features.append(np.mean(librosa.feature.mfcc(x, sr=sample_rate, n_mfcc=40).T,axis=0))
    labels.append(train.Class[i])

# In[6]:

#extracting features from test set
features_test=[]
for i in range(len(test)):
    test_dir = "D:\\Urban Sound Classification\\test\\Test\\"
    filename=os.path.join(os.path.abspath(test_dir)+'\\'+str(test.ID[i])+'.wav')
    x, sample_rate = librosa.load(filename, res_type='kaiser_fast')
    features_test.append(np.mean(librosa.feature.mfcc(x, sr=sample_rate, n_mfcc=40).T,axis=0))

# In[7]:

#label encoding the categorical data (labels)
#making test and train

x=np.array(features)
x_test=np.array(features_test)

#encoding the labels
from sklearn.preprocessing import LabelEncoder
from keras.utils import to_categorical

le = LabelEncoder()
y = to_categorical(le.fit_transform(labels))

# In[8]:

print('Shape of Features(Train): ',x.shape)
print('Shape of Features(Test): ',x_test.shape)
print('Shape of Labels(Train): ',y.shape)

# # Nueral Network

# In[9]:

import tensorflow as tf

model1=tf.keras.models.Sequential([

    tf.keras.layers.Dense(256,activation='relu',input_shape=(40,)),
    #using dropout regularization to randomly drop nodes , reduce overfitting, and improve generalization
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(256,activation='relu'),
    tf.keras.layers.Dropout(0.5),

    #10 nodes since we have 10 labels
    tf.keras.layers.Dense(10,activation='softmax')
])

# In[10]:

#using tensorflow callbacks to reduce the learning rate if loss does not improve

reduce =tf. keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10, mode='auto')

model1.compile(optimizer='adam',loss='categorical_crossentropy',metrics=['accuracy'])

history1=model1.fit(x,y, batch_size=32, epochs=100, validation_split=0.1,callbacks=[reduce])

# In[11]:

#plot of training and validation accuracy/loss per epoch

acc=history1.history['accuracy']
val_acc=history1.history['val_accuracy']
loss=history1.history['loss']
val_loss=history1.history['val_loss']

epochs=range(len(acc))

plt.plot(epochs,acc,'b',label='Training Accuracy')
plt.plot(epochs,val_acc,'r',label='Validation Accuracy')
plt.legend()
plt.figure()

plt.plot(epochs,loss,'b',label='Training Loss')
plt.plot(epochs,val_loss,'r',label='Validation Loss')
plt.legend()
plt.show()

# In[13]:

nn_pred=model1.predict(x_test)
nn_max=np.argmax(nn_pred,axis=1)
nn_pred_test=le.inverse_transform(nn_max)

# In[14]:

proba_nn = model1.predict_proba(x_test).ravel()

# In[22]:

nn_test_label = pd.DataFrame({'ID':test['ID'].values})
nn_test_label.insert(loc=0, column='Class', value=nn_pred_test)

# In[24]:

nn_test_label.head()

# # CNN

# In[25]:

max_pad_len=174 #defining a maximum length for each column

file_name = "D:\\Urban Sound Classification\\train\\Train\\8279.wav"
audio, sample_rate = librosa.load(file_name, res_type='kaiser_fast')
mfcc=librosa.feature.mfcc(y=audio, sr=sample_rate, n_mfcc=40)
print('Shape of feature matrix: ',mfcc.shape)
pad_width = max_pad_len - mfcc.shape[1]
print('Pad width: ',pad_width)
mfcc = np.pad(mfcc, pad_width=((0, 0), (0, pad_width)), mode='constant')
print('Shape of feature matrix after padding: ',mfcc.shape)

# In[26]:

#CNN does not support varying column length.fix by setting max length and padding
features_padded=[]
labels=[]
pad=174 # defining max column length as 174 for variation

for i in range(len(train)):
    file_name = "D:\\Urban Sound Classification\\train\\Train\\"
    filename=os.path.join(os.path.abspath(file_name)+'\\'+str(train.ID[i]) +'.wav')
    x, sample_rate = librosa.load(filename, res_type='kaiser_fast')
    mfcc=librosa.feature.mfcc(y=x, sr=sample_rate, n_mfcc=40)
    padding = pad - mfcc.shape[1]
    mfcc = np.pad(mfcc, pad_width=((0, 0), (0, padding)), mode='constant')
    features_padded.append(mfcc)
    labels.append(train.Class[i])

# In[27]:

features_test_padded=[]

for i in range(len(test)):
    file_name = "D:\\Urban Sound Classification\\test\\Test\\"
    filename=os.path.join(os.path.abspath(file_name)+'\\'+str(test.ID[i])+'.wav')
    x, sample_rate = librosa.load(filename, res_type='kaiser_fast')
    mfcc=librosa.feature.mfcc(y=x, sr=sample_rate, n_mfcc=40)
    padding = pad - mfcc.shape[1]
    mfcc = np.pad(mfcc, pad_width=((0, 0), (0, padding)), mode='constant')
    features_test_padded.append(mfcc)

# In[28]:

x_new=np.array(features_padded)
x_new=x_new.reshape(x_new.shape[0], 40, 174, 1) #reshaping the array for the convolution layer input

x_new_test=np.array(features_test_padded)
x_new_test=x_new_test.reshape(x_new_test.shape[0], 40, 174, 1)

#encoding the labels
from sklearn.preprocessing import LabelEncoder
from keras.utils import to_categorical

le = LabelEncoder()
y = to_categorical(le.fit_transform(labels))

# In[29]:

print('Shape of Features(Train): ',x_new.shape)
print('Shape of Features(Test): ',x_new_test.shape)
print('Shape of Labels(Train): ',y.shape)

# In[30]:

import tensorflow as tf

model2=tf.keras.models.Sequential([

    tf.keras.layers.Conv2D(16,(3,3),activation='relu',input_shape=(40,174,1)),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Conv2D(32,(3,3),activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Conv2D(64,(3,3),activation='relu'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Dropout(0.2),
    tf.keras.layers.Conv2D(128,(3,3),activation='relu',padding='same'),
    tf.keras.layers.MaxPooling2D(2,2),
    tf.keras.layers.Dropout(0.4),

    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(256,activation='relu'),
    tf.keras.layers.Dropout(0.5),
    tf.keras.layers.Dense(256,activation='relu'),
    tf.keras.layers.Dropout(0.5),

    tf.keras.layers.Dense(10,activation='softmax')
])

# In[31]:

reduce =tf. keras.callbacks.ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10, mode='auto')

model2.compile(optimizer='adam',loss='categorical_crossentropy',metrics=['accuracy'])
history2=model2.fit(x_new,y, batch_size=32, epochs=100, validation_split=0.1,callbacks=[reduce])

# In[32]:

acc=history2.history['accuracy']
val_acc=history2.history['val_accuracy']
loss=history2.history['loss']
val_loss=history2.history['val_loss']

epochs=range(len(acc)) #No. of epochs

#Plot training and validation accuracy per epoch
import matplotlib.pyplot as plt
plt.plot(epochs,acc,'b',label='Training Accuracy')
plt.plot(epochs,val_acc,'r',label='Validation Accuracy')
plt.legend()
plt.figure()

#Plot training and validation loss per epoch
plt.plot(epochs,loss,'b',label='Training Loss')
plt.plot(epochs,val_loss,'r',label='Validation Loss')
plt.legend()
plt.show()

# In[44]:

cnn_pred=model2.predict(x_new_test)
cnn_max=np.argmax(cnn_pred,axis=1)
cnn_pred_test=le.inverse_transform(cnn_max)

# In[45]:

proba_cnn = model2.predict_proba(x_new_test).ravel()

# In[47]:

cnn_test_label= pd.DataFrame({'ID':test['ID'].values})
cnn_test_label.insert(loc=0, column='Class', value=cnn_pred_test)

# In[48]:

cnn_test_label.head()


from django import template

register = template.Library()

@register.filter(name="exist")
def exist(list):
    if list and len(list) >= 1:
        return all(x is not None for x in list[0])
    else:
        return False


import argparse
import numpy as np
import json
import re
import random
import itertools
import torch
import torch.nn as nn
import torch.nn.functional as F

import data
import models
import utils

from tqdm import tqdm
from pathlib import Path
from torch.utils.data import DataLoader

def get_metrics(recalls):
    metrics = {}
    for name, value in recalls.items():
        if name == 'MR':
            metrics[name] = np.median(value)
        else:
            metrics[f'R@{name}'] = np.mean(value)*100
    return metrics

def evaluate(model, video_iterator, lang_iterator, annotations, device, preliminary=100,
                model_types=['model'], iou_thresholds=[0.5,0.7]):

    videos = {}
    # compute embeddings for all clips in each unique video in dataset
    for batch in tqdm(video_iterator):
        with torch.no_grad():
            videos[batch['video']] = model(batch['feature'].to(device))

    recalls = {(model_type, iou_thr): {1: [], 10: [], 100: [], 'MR':[]}
        for (model_type, iou_thr) in itertools.product(model_types, iou_thresholds)}
    moments = lang_iterator.batch_sampler.moments
    print('\nEvaluation:')
    # compute embeddings for all queries in dataset
    for li, batch in tqdm(enumerate(lang_iterator)):
        with torch.no_grad():
            lang_emb = model(batch['feature'].to(device), False, device)

        # ground truth: 1 if IoU between moment and at least 2 manual annotations is greater the threshold
        # for incorrect videos all IoU score are 0
        distances, predicts = [], {key: [] for key in recalls.keys()}
        for video in videos.keys():
            video_emb = videos[video]
            num_segments = video_emb.size(0)
            # for each query compute distance between query and clip embeddings
            dist = F.pairwise_distance(video_emb, lang_emb.repeat(num_segments, 1))
            # for each possible moment in video compute distance between moment and query
            # as sum of distances between clips in the moment and query (precomputed on prev step)
            # (15 possible moments for video of 5 clips, 21 - for 6 clips)
            for start_t, end_t in moments[num_segments]:
                distances.append(dist.index_select(0, torch.arange(start_t, end_t+1).to(device)).mean().item())
                if video == batch['video']:
                    ious = utils.get_iou(annotations[batch['annot_id']]['times'], start_t, end_t)
                    for iou_thr in iou_thresholds:
                        predicts[('model', iou_thr)].append(int(( ious > iou_thr ).sum() >= 2))
                else:
                    for iou_thr in iou_thresholds:
                        predicts[('model', iou_thr)].append(0)

        # sort all moments from all videos by distances
        ind_rand = np.random.choice(np.arange(len(distances)), size=len(distances), replace=False)
        for iou_thr in iou_thresholds:
            ground_truth = np.array(predicts[('model',iou_thr)], dtype=int)
            predicts[('model', iou_thr)] = ground_truth[np.argsort(distances)]
            predicts[('chance',iou_thr)] = ground_truth[ind_rand]
            for model_type in model_types:
                for k in recalls[(model_type, iou_thr)].keys():
                    if k == 'MR':
                        # median rank: index position of first correct moment
                        recalls[(model_type, iou_thr)][k].append(np.where(predicts[(model_type, iou_thr)] == 1)[0][0])
                    else:
                        # Recall@K: 1 if correct moment is among in top-K moments, 0 otherwise
                        recalls[(model_type, iou_thr)][k].append(int(predicts[(model_type, iou_thr)][:k].sum() > 0))

        if (li % preliminary == 0) and (li != 0):
            print()
            for (model_type, iou_thr) in recalls.keys():
                print(f'{model_type}, IoU={iou_thr}:\t', ''.join([f'{name}: {value:.4f}\t'
                    for name, value in get_metrics(recalls[(model_type, iou_thr)]).items()]))
            print()

    return {f'{model_type}, IoU={iou_thr}': get_metrics(recalls[(model_type, iou_thr)])
                            for (model_type, iou_thr) in recalls.keys()}

if __name__ == '__main__':

    torch.random.manual_seed(123)
    random.seed(123)
    np.random.seed(123)

    parser = argparse.ArgumentParser()
    parser.add_argument("--features_dir", type=str, default='features/')
    parser.add_argument("--embedding_dir", type=str, default='glove_pretrained/')
    parser.add_argument("--dataset_dir", type=str, default='dataset/')
    parser.add_argument("--experiment_dir", type=str, default='experiment/')
    parser.add_argument("--missed_videos", type=str, default='../')
    parser.add_argument("--device", type=str, default=None)
    parser.add_argument("--validation", type=str, default='val')
    parser.add_argument("--preliminary_print", type=int, default=100)
    parser.add_argument("--iou", type=float, default=0.5)
    parser.add_argument("--model_types", type=str, nargs='+', default=[])
    parser.add_argument("--normalize_lang", type=utils.str2bool, default=False)
    args = parser.parse_args()

    FT_DIRECTORY = args.features_dir
    EMB_DIRECTORY = args.embedding_dir
    DATASET_DIRECTORY = args.dataset_dir
    EXPER_DIRECTORY = args.experiment_dir
    DATASET_TYPE = args.validation

    if args.device in ['cpu','cuda']:
        DEVICE = args.device
    else:
        if torch.cuda.is_available():
            DEVICE = torch.device('cuda')
        else:
            DEVICE = torch.device('cpu')

    print('Loading word embeddings:')
    word_indexer = data.WordIndexer(EMB_DIRECTORY)

    # load state dict
    state = torch.load(Path(EXPER_DIRECTORY).joinpath('last.pth'))
    FEATURE_TYPE = state['params']['feature_type']
    POOLING = state['params']['pooling']
    normalize_lang = state['params']['normalize_lang']
    normalize_loss = state['params']['normalize_loss']
    MODEL_TYPES = [model_type for model_type in args.model_types
                            if model_type in ['model','chance']]
    if MODEL_TYPES == []:
        MODEL_TYPES = ['model']
    print(DEVICE, FEATURE_TYPE, DATASET_TYPE, args.iou, MODEL_TYPES)

    moment_frequency = {}
    if 'prior' in MODEL_TYPES:

        print('Loading train dataset:')
        train_annotations, train_videos = utils.load_dataset_info('train', DATASET_DIRECTORY, args.missed_videos)
        train_dataset = data.CustomDataset(
            train_videos, train_annotations, word_indexer, FT_DIRECTORY, FEATURE_TYPE, pooling=POOLING)

        for _, annot_info in train_annotations.items():
            num_segments = train_dataset.num_segments_info[annot_info['video']]
            if num_segments not in moment_frequency.keys():
                moment_frequency[num_segments] = {tuple(moment): 0 for moment in utils.generate_moments(num_segments)}
                moment_frequency[num_segments]['total'] = 0
            else:
                for t in annot_info['times']:
                    moment_frequency[num_segments][tuple(t)] += 1
                    moment_frequency[num_segments]['total'] += 1

        for num_segments in moment_frequency.keys():
            moment_frequency[num_segments] = [freq / moment_frequency[num_segments]['total']
                        for moment, freq in moment_frequency[num_segments].items() if moment != 'total']

    print('Loading validation dataset:')
    val_annotations, val_videos = utils.load_dataset_info(DATASET_TYPE, DATASET_DIRECTORY, args.missed_videos)

    val_dataset = data.CustomDataset(val_videos, val_annotations, word_indexer, FT_DIRECTORY, FEATURE_TYPE, validate=True)
    val_video_iter = DataLoader(val_dataset, shuffle=False, collate_fn=data.validate_collate,
        batch_sampler=data.VideoBatchSampler(val_videos, val_dataset.num_segments_info))
    val_lang_iter = DataLoader(val_dataset, shuffle=False, collate_fn=data.validate_collate,
        batch_sampler=data.LanguageBatchSampler(val_annotations, val_dataset.num_segments_info))

    model = models.CALModel(
        pretrained_emb=word_indexer.get_embeddings(),
        visual_input_dim=data.FEATURE_DIM[FEATURE_TYPE]*2+2,
        emb_dim=data.EMBEDDING_DIM,
        normalize_lang=args.normalize_lang
    )
    model.load_state_dict(state['model_state_dict'])
    model = model.to(DEVICE)
    model.eval()

    # evaluate model with exhaustive search
    metrics = evaluate(
        model, val_video_iter, val_lang_iter, val_annotations, DEVICE, args.preliminary_print, MODEL_TYPES)

    for model_type in metrics.keys():
        print(f'{model_type}\t',''.join([f'{name}: {value:.4f}\t'
            for name, value in metrics[model_type].items()]))

    with open(Path(EXPER_DIRECTORY).joinpath('metrics.json'), 'w') as f:
        json.dump(metrics, f)


__author__ = "Kellerberrin"

package = "ProductAuthentication"

"""Product Authentication using cryptographic signatures and X509 certificates."""

import endpoints
from protorpc import messages
from protorpc import message_types
from protorpc import remote

from authenticate import verifySerialCode

class RequestVerifySerial(messages.Message):
    """Extract the drug database search arguments."""
    publickey = messages.StringField(1, required=True)
    response = messages.StringField(2, required=True)
    barcode = messages.StringField(3, required=True)
    serialcode = messages.StringField(4, required=True)
    ipaddress = messages.StringField(5, required=True)

class VerifySerialResult(messages.Message):
    """Return the json object"""
    resultMessage = messages.StringField(1)

@endpoints.api(name="productAuth", version="v1")
class ProductAuthentication(remote.Service):
    """Product Authentication API v1."""

    REQUEST_VERIFYSERIAL_RESOURCE = endpoints.ResourceContainer(
        RequestVerifySerial,
        publickey=messages.StringField(1, required=True),
        response=messages.StringField(2, required=True),
        barcode=messages.StringField(3, required=True),
        serialcode=messages.StringField(4, required=True),
        ipaddress=messages.StringField(5, required=True),

    )

    @endpoints.method(REQUEST_VERIFYSERIAL_RESOURCE
                      , VerifySerialResult
                      , path="verifySerial"
                      , http_method="POST"
                      , name="verifySerial")
    def verifyserial_implementation(self, request):
        return VerifySerialResult(resultMessage=verifySerialCode(request.publickey,
                                                                 request.response,
                                                                 request.barcode,
                                                                 request.serialcode,
                                                                 request.ipaddress))

application = endpoints.api_server([ProductAuthentication],  restricted=False)


import logging

from django.apps import apps
from django.core import serializers
from django.core.exceptions import ObjectDoesNotExist, MultipleObjectsReturned
from django.shortcuts import get_object_or_404, Http404
from django.db.models import Count, Q
from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework.decorators import action
from rest_framework.exceptions import AuthenticationFailed, MethodNotAllowed
from rest_framework.viewsets import ModelViewSet

from surf.apps.core.search import get_search_client
from surf.apps.communities.models import Team, Community
from surf.apps.materials.models import (
    Collection,
    Material,
    CollectionMaterial,
    SharedResourceCounter,
    RESOURCE_TYPE_COLLECTION
)
from surf.apps.materials.serializers import (
    CollectionSerializer,
    CollectionMaterialsRequestSerializer,
    MaterialShortSerializer,
    CollectionMaterialPositionSerializer,
)
from surf.apps.materials.utils import add_extra_parameters_to_materials, get_material_details_by_id
from surf.apps.materials.filters import CollectionFilter

logger = logging.getLogger(__name__)
filters_app = apps.get_app_config("filters")

class CollectionMaterialPromotionAPIView(APIView):

    def post(self, request, *args, **kwargs):
        # only active and authorized users can promote materials in the collection
        collection_instance = Collection.objects.get(id=kwargs['collection_id'])

        # check whether the material is actually in the collection
        external_id = kwargs['external_id']
        collection_materials = CollectionMaterial.objects.filter(
            collection=collection_instance, material__external_id=external_id)
        if not collection_materials:
            raise Http404()

        # The material should only be in the collection once
        assert len(collection_materials) == 1, f"Material with id {external_id} is in collection " \
                                               f"{collection_instance} multiple times."
        collection_material = collection_materials[0]
        # promote or demote the material
        collection_material.featured = not collection_material.featured
        collection_material.save()

        return Response(serializers.serialize('json', [collection_material]))

class CollectionViewSet(ModelViewSet):
    """
    View class that provides CRUD methods for Collection and `get`, `add`
    and `delete` methods for its materials.
    """

    queryset = Collection.objects \
        .filter(deleted_at=None) \
        .annotate(community_cnt=Count('communities', filter=Q(deleted_at=None)))
    serializer_class = CollectionSerializer
    filter_class = CollectionFilter
    permission_classes = []

    def retrieve(self, request, *args, **kwargs):
        instance = self.get_object()

        shared = request.GET.get("shared")
        if shared:
            # increase sharing counter
            counter_key = SharedResourceCounter.create_counter_key(
                RESOURCE_TYPE_COLLECTION,
                str(instance.id),
                share_type=shared)

            SharedResourceCounter.increase_counter(counter_key, extra=shared)

        serializer = self.get_serializer(instance)
        return Response(serializer.data)

    def get_object(self):
        obj = get_object_or_404(self.get_queryset(), pk=self.kwargs["pk"])
        self.check_object_permissions(self.request, obj)
        if self.request.method != 'GET':
            check_access_to_collection(self.request.user, obj)
        return obj

    @action(methods=['get', 'post', 'delete'], detail=True)
    def materials(self, request, pk=None, **kwargs):
        instance = self.get_object()

        if request.method == "GET":
            # validate request parameters
            serializer = CollectionMaterialsRequestSerializer(data=request.GET)
            serializer.is_valid(raise_exception=True)
            data = serializer.validated_data
            ids = [m.external_id for m in instance.materials.order_by("id").filter()]

            rv = dict(records=[],
                      records_total=0,
                      filters=[],
                      page=data["page"],
                      page_size=data["page_size"])

            if ids:
                client = get_search_client()

                res = client.get_materials_by_id(ids, 1, len(ids))
                records = res.get("records", [])
                records = add_extra_parameters_to_materials(filters_app.metadata, records)

                collection_materials = CollectionMaterial.objects.filter(
                    collection=instance, material__external_id__in=[r['external_id'] for r in records]
                )

                for collection_material in collection_materials:
                    record = next(r for r in records if r['external_id'] == collection_material.material.external_id)
                    record['featured'] = collection_material.featured
                    record['position'] = collection_material.position

                rv["records"] = records
                rv["records_total"] = res["recordcount"]

            return Response(rv)

        data = []
        for d in request.data:
            # validate request parameters
            if request.method == "POST":
                serializer = CollectionMaterialPositionSerializer(data=d)
            elif request.method == "DELETE":
                serializer = MaterialShortSerializer(data=d)
            else:
                raise MethodNotAllowed(request.method, detail="Method not supported")

            serializer.is_valid(raise_exception=True)
            data.append(serializer.validated_data)

        if request.method == "POST":
            self._add_materials(instance, data)

        elif request.method == "DELETE":
            self._delete_materials(instance, data)

        res = MaterialShortSerializer(many=True).to_representation(
            instance.materials.filter(deleted_at=None)
        )
        return Response(res)

    @staticmethod
    def _add_materials(instance, materials):
        """
        Add materials to collection
        :param instance: collection instance
        :param materials: added materials
        :return:
        """

        for material in materials:
            m_external_id = material["external_id"]
            m_position = material["position"]

            details = get_material_details_by_id(m_external_id)
            if not details:
                continue

            m, _ = Material.objects.update_or_create(external_id=m_external_id)
            CollectionMaterial.objects.create(collection=instance, material=m, position=m_position)

    @staticmethod
    def _delete_materials(instance, materials):
        """
        Delete materials from collection
        :param instance: collection instance
        :param materials: materials that should be removed from collection
        :return:
        """

        materials = [m["external_id"] for m in materials]
        materials = Material.objects.filter(external_id__in=materials).all()
        instance.materials.remove(*materials)

def check_access_to_collection(user, instance=None):
    """
    Check if user is active and owner of collection (if collection
    is not None)
    :param user: user
    :param instance: collection instance
    :return:
    """
    if not user or not user.is_authenticated:
        raise AuthenticationFailed()
    try:
        community = Community.objects.get(collections__in=[instance])
        Team.objects.get(community=community, user=user)
    except ObjectDoesNotExist as exc:
        raise AuthenticationFailed(f"User {user} is not a member of a community that has collection {instance}. "
                                   f"Error: \"{exc}\"")
    except MultipleObjectsReturned as exc:
        logger.warning(f"The collection {instance} is in multiple communities. Error:\"{exc}\"")
        communities = Community.objects.filter(collections__in=[instance])
        teams = Team.objects.filter(community__in=communities, user=user)
        if len(teams) > 0:
            logger.debug("At least one team satisfies the requirement of be able to delete this collection.")
        else:
            raise AuthenticationFailed(f"User {user} is not a member of any community with collection {instance}. "
                                       f"Error: \"{exc}\"")


#!/usr/bin/env python3  # ~~~ VERIFIED ~~~ #

# region         === importations ===         NOIGER #
import math
import random

from typing import Union
from typing import Tuple
from typing import Optional
from typing import List
from dataclasses import dataclass

import pandas as pd

import torch
from torch import nn
from torch import Tensor
from torch import FloatTensor
from torch import LongTensor
from torch.utils.data import DataLoader
from torch.optim import AdamW
from torch.nn.modules.loss import L1Loss
from torch.nn import CrossEntropyLoss
from torch.nn import MSELoss
from torch.nn import CTCLoss

from transformers import BartModel
from transformers import BartForConditionalGeneration
from transformers import BartConfig
from transformers import get_linear_schedule_with_warmup
from transformers.modeling_outputs import BaseModelOutput
from transformers.modeling_outputs import Seq2SeqModelOutput
from transformers.modeling_outputs import Seq2SeqLMOutput
from transformers.modeling_outputs import CausalLMOutput
from transformers.generation_utils import GenerationMixin
from transformers import logging
from transformers.models.bart.modeling_bart import BartEncoder
from transformers.models.bart.modeling_bart import BartDecoder
from transformers.models.bart.modeling_bart import shift_tokens_right
from transformers.models.bart.modeling_bart import BartLearnedPositionalEmbedding, BartEncoderLayer
from transformers.models.bart.modeling_bart import _expand_mask

import pytorch_lightning as pl

try:
    from .torch_cif import cif_function
    from .utils import mask_generator
    from .datasets import DataSetCollectorGeneral
except:
    from torch_cif import cif_function
    from utils import mask_generator
    from datasets import DataSetCollectorGeneral
# endregion      === importations ===      NOIGERDNE #

# region       === aug_dataclasses ===        NOIGER #
@dataclass
class AugBaseModelOutput(BaseModelOutput):
    last_hidden_state: torch.FloatTensor = None
    out_attention_mask: Optional[torch.FloatTensor] = None
    length_loss: Optional[Tuple[torch.FloatTensor]] = None
    pred_word_lengths: Optional[torch.LongTensor] = None
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None

@dataclass
class AugSeq2SeqModelOutput(Seq2SeqModelOutput):
    last_hidden_state: torch.FloatTensor = None
    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
    decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
    cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
    encoder_last_hidden_state: Optional[torch.FloatTensor] = None
    encoder_last_hidden_out_attention_mask: Optional[torch.FloatTensor] = None
    encoder_length_loss: Optional[Tuple[torch.FloatTensor]] = None
    encoder_pred_word_lengths: Optional[torch.LongTensor] = None
    encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None

@dataclass
class AugSeq2SeqLMOutput(Seq2SeqLMOutput):
    loss: Optional[torch.FloatTensor] = None
    logits: torch.FloatTensor = None
    past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None
    decoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    decoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
    cross_attentions: Optional[Tuple[torch.FloatTensor]] = None
    encoder_last_hidden_state: Optional[torch.FloatTensor] = None
    encoder_last_hidden_out_attention_mask: Optional[torch.FloatTensor] = None
    encoder_length_loss: Optional[Tuple[torch.FloatTensor]] = None
    encoder_pred_word_lengths: Optional[torch.LongTensor] = None
    encoder_hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    encoder_attentions: Optional[Tuple[torch.FloatTensor]] = None
    masked_lm_loss: Optional[torch.FloatTensor] = None
    real_length_loss: Optional[torch.FloatTensor] = None
# endregion    === aug_dataclasses ===     NOIGERDNE #

logger = logging.get_logger("transformers.models.bart.modeling_bart")

class SentBartEncoder(BartEncoder):
    def __init__(self, config: BartConfig, embed_tokens: Optional[nn.Embedding] = None):
        super().__init__(config, embed_tokens)
        embed_dim = config.d_model
        self.collapse_n = getattr(config, "collapse_n", 0)
        self.config.collapse_n = self.collapse_n
        if self.collapse_n == -1:
            self.skip_cif = True
        else:
            self.skip_cif = False
        self.use_self = getattr(config, "use_self", False)
        self.config.use_self = self.use_self

        self.post_initialization(embed_dim, word_extractor=cif_function)

    def post_initialization(self, embed_dim, word_extractor=cif_function):
        # == encoder == #
        self.alpha_predictor = nn.Linear(embed_dim, 1)  # TODO: check!
        self.word_extractor = word_extractor
        self.length_predictor = nn.Linear(embed_dim, 1)

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        word_length_tensor: Optional[torch.LongTensor] = None,
        alpha_values: Optional[torch.FloatTensor] = None,
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
      ) -> Union[Tuple, AugBaseModelOutput]:
        encoder_outputs = super().forward(
            input_ids,
            attention_mask,
            head_mask,
            inputs_embeds,
            output_attentions,
            output_hidden_states,
            return_dict,
        )
        encoder_outputs = BaseModelOutput(
            last_hidden_state=encoder_outputs[0],
            hidden_states=encoder_outputs[1]
                if len(encoder_outputs) > 1 else None,
            attentions=encoder_outputs[2]
                if len(encoder_outputs) > 2 else None,
        )

        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        if self.skip_cif:
            hidden_states = encoder_outputs.last_hidden_state
            out_attention_mask = attention_mask
            length_loss = (
                torch.zeros(len(hidden_states)).float(),
                torch.zeros(len(hidden_states)).long(),)
            pred_word_lengths = None
        else:
            if self.collapse_n == 1:    # 1 for all
                word_length_tensor = torch.ones_like(word_length_tensor)
            elif self.collapse_n == 2:  # all for all
                word_length_tensor = (attention_mask.sum(-1).long()
                                      .to(word_length_tensor.device))
                                        # else: N for all
            (hidden_states, out_attention_mask, length_loss,
                          # out_attention_mask := (
                          #     encoder_output_attention_mask)
             pred_word_lengths,
             _, _) = self.sent_retriever(
                encoder_outputs.last_hidden_state,
                word_length_tensor=word_length_tensor,
                alpha_values=alpha_values,
                padding_mask=1 - attention_mask if attention_mask is not None else None)
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

        if not return_dict:
            return tuple(v for v in [
                hidden_states,
                out_attention_mask,
                encoder_outputs.hidden_states,
                length_loss,
                pred_word_lengths,
                encoder_outputs.attentions,
            # ] if v is not None)
            ] if v is not None)
        return AugBaseModelOutput(
            last_hidden_state=hidden_states,
            out_attention_mask=out_attention_mask,
            length_loss=length_loss,
            pred_word_lengths=pred_word_lengths,
            hidden_states=encoder_outputs.hidden_states,
            # # attentions=encoder_outputs.attentions,
            # attentions=encoder_outputs.attentions,
        )

    def sent_retriever(self,
        encoder__last_hidden_state,
        word_length_tensor=None,
        alpha_values=None,
        padding_mask=None,
        return_all=False,
        return_original=False,
      ):
        if alpha_values is None:
            alpha_values = self.alpha_predictor(encoder__last_hidden_state)
            alpha_values = alpha_values.squeeze(-1).sigmoid()  # B x S

        if word_length_tensor is None or self.use_self:
            # print("No given! self predict")
            word_length_tensor = alpha_values.sum(-1).long()
        else:
            # print("Wordlen given")
            # predicted_word_length_tensor = alpha_values.sum(-1).long()
            pass

        encoder__word_representations_CIF = (
            self.word_extractor(
                encoder__last_hidden_state,
                alpha=alpha_values,
                padding_mask=padding_mask,
                target_lengths=word_length_tensor,
            )
        )
        # TODO: Keep all CIF
        [encoder_word_representation] = encoder__word_representations_CIF['cif_out']
        [pred_word_lengths] = encoder__word_representations_CIF['alpha_sum']
        encoder_word_representation = encoder_word_representation.contiguous()
        # pred_word_lengths = pred_word_lengths.contiguous()
        # length_loss = 0.
        length_loss = ((pred_word_lengths, word_length_tensor,)
            if word_length_tensor is not None else
            (pred_word_lengths, pred_word_lengths,))
            # aliased as `encoder_word_representation`
            # FIXME: distributed problem!
            # TODO: add other CIF ouptuts!
        # length_loss = length_loss.contiguous()

        encoder_output_attention_mask = (
            # mask_generator(word_length_tensor)
            mask_generator(word_length_tensor)
            if word_length_tensor is not None else
            mask_generator(pred_word_lengths))
            # TODO: check length prediction!

        return (
            encoder_word_representation,
            encoder_output_attention_mask,
            length_loss,
            pred_word_lengths,
            encoder__word_representations_CIF if return_all else None,
            # encoder__last_hidden_state if return_original else None,
            encoder__last_hidden_state if return_original else None,
        )

class SentBart(BartModel):
    def __init__(self, config: BartConfig):
        super(BartModel, self).__init__(config)

        padding_idx, vocab_size = config.pad_token_id, config.vocab_size
        self.shared = nn.Embedding(vocab_size, config.d_model, padding_idx)

        self.encoder = SentBartEncoder(config, self.shared)
        self.decoder = BartDecoder(config, self.shared)

        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        word_length_tensor: Optional[torch.LongTensor] = None,
        alpha_values: Optional[torch.FloatTensor] = None,
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        decoder_head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[List[torch.FloatTensor]] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, AugSeq2SeqModelOutput]:

        # different to other models, Bart automatically creates decoder_input_ids from
        # input_ids if no decoder_input_ids are provided
        if decoder_input_ids is None and decoder_inputs_embeds is None:
            if input_ids is None:
                raise ValueError(
                    "If no `decoder_input_ids` or `decoder_inputs_embeds` are "
                    "passed, `input_ids` cannot be `None`. Please pass either "
                    "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`."
                )

            decoder_input_ids = shift_tokens_right(
                input_ids, self.config.pad_token_id, self.config.decoder_start_token_id
            )

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if encoder_outputs is None:
            encoder_outputs = self.encoder(
                input_ids=input_ids,
                attention_mask=attention_mask,
                # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
                word_length_tensor=word_length_tensor,
                alpha_values=alpha_values,
                # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
                head_mask=head_mask,
                inputs_embeds=inputs_embeds,
                output_attentions=output_attentions,
                output_hidden_states=output_hidden_states,
                return_dict=return_dict,
            )
        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        # If the user passed a tuple for encoder_outputs, we wrap it in an AugBaseModelOutput when return_dict=True
        elif return_dict and not isinstance(encoder_outputs, AugBaseModelOutput):
            encoder_outputs = AugBaseModelOutput(
                last_hidden_state=encoder_outputs[0],
                out_attention_mask=encoder_outputs[1],
                length_loss=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
                pred_word_lengths=encoder_outputs[3] if len(encoder_outputs) > 3 else None,
                hidden_states=encoder_outputs[4] if len(encoder_outputs) > 4 else None,
                # attentions=encoder_outputs[5] if len(encoder_outputs) > 5 else None,
                    attentions=encoder_outputs[5] if len(encoder_outputs) > 5 else None,
            )
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

        # decoder outputs consists of (dec_features, past_key_value, dec_hidden, dec_attn)
        decoder_outputs = self.decoder(
            input_ids=decoder_input_ids,
            attention_mask=decoder_attention_mask,
            encoder_hidden_states=encoder_outputs[0],
            # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
            encoder_attention_mask=encoder_outputs[1],
            # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
            head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        if not return_dict:
            return decoder_outputs + encoder_outputs

        return AugSeq2SeqModelOutput(
            last_hidden_state=decoder_outputs.last_hidden_state,
            past_key_values=decoder_outputs.past_key_values,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
            cross_attentions=decoder_outputs.cross_attentions,
            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
            encoder_last_hidden_out_attention_mask=encoder_outputs.out_attention_mask,
            encoder_length_loss=encoder_outputs.length_loss,
            encoder_pred_word_lengths=encoder_outputs.pred_word_lengths,
            encoder_hidden_states=encoder_outputs.hidden_states,
            # # encoder_attentions=encoder_outputs.attentions,
            # encoder_attentions=encoder_outputs.attentions,
        )

    # region CHANGED >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    def fix_encoder_(self, to_fix=True):
        self.encoder.requires_grad_(not to_fix)
    # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

class SentBartForConditionalGeneration(BartForConditionalGeneration):
    def __init__(self, config: BartConfig):
        super(BartForConditionalGeneration, self).__init__(config)
        self.model = SentBart(config)
        self.register_buffer(
            "final_logits_bias",
            torch.zeros((1, self.model.shared.num_embeddings)))
        self.lm_head = nn.Linear(
            config.d_model,
            self.model.shared.num_embeddings,
            bias=False)

        self.weight_len = getattr(config, "weight_len", None)
        self.config.weight_len = self.weight_len
        self.minimize_len = getattr(config, "minimize_len", False)
        self.config.minimize_len = self.minimize_len

        # Initialize weights and apply final processing
        self.post_init()

    # region CHANGED >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
    def fix_encoder_(self, to_fix=True):
        self.model.fix_encoder_(to_fix)
    # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        # region CHANGED >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        word_length_tensor: Optional[torch.LongTensor] = None,
        alpha_values: Optional[torch.FloatTensor] = None,
        # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        decoder_input_ids: Optional[torch.LongTensor] = None,
        decoder_attention_mask: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        decoder_head_mask: Optional[torch.Tensor] = None,
        cross_attn_head_mask: Optional[torch.Tensor] = None,
        encoder_outputs: Optional[List[torch.FloatTensor]] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        decoder_inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, AugSeq2SeqLMOutput]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

        Returns:
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if labels is not None:
            if use_cache:
                logger.warning("The `use_cache` argument is changed to `False` since `labels` is provided.")
            use_cache = False
            if decoder_input_ids is None and decoder_inputs_embeds is None:
                decoder_input_ids = shift_tokens_right(
                    labels, self.config.pad_token_id, self.config.decoder_start_token_id
                )

        # print(self.model.encoder.layers[5].fc1.weight[0][:6] * 100)
        # print(self.model.encoder.alpha_predictor.weight[0][:6] * 100)
        # print(self.model.decoder.layers[0].fc1.weight[0][:6] * 100)

        if self.minimize_len:
            word_length_tensor = None

        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
            # region CHANGED >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
            word_length_tensor=word_length_tensor,
            alpha_values=alpha_values,
            # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
            decoder_input_ids=decoder_input_ids,
            encoder_outputs=encoder_outputs,
            decoder_attention_mask=decoder_attention_mask,
            head_mask=head_mask,
            decoder_head_mask=decoder_head_mask,
            cross_attn_head_mask=cross_attn_head_mask,
            past_key_values=past_key_values,
            inputs_embeds=inputs_embeds,
            decoder_inputs_embeds=decoder_inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        lm_logits = self.lm_head(outputs[0]) + self.final_logits_bias

        masked_lm_loss = None
        if labels is not None:
            # region CHANGED >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
            loss_fct = CrossEntropyLoss()
            masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1))
            # loss_fct = nn.CTCLoss(1)
            # masked_lm_loss = loss_fct(
            #     lm_logits, # $$$
            #     labels,
            #     attention_mask.sum(-1),
            #     (labels != 1).sum(-1))
            # breakpoint()
            # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

        real_length_loss = None
        if word_length_tensor is not None and self.weight_len is not None:
            # region CHANGED >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
            loss_fct = nn.L1Loss()
            # real_length_loss = torch.linalg.vector_norm(outputs.encoder_length_loss.float(), ord=1)
            real_length_loss = loss_fct(*(outputs.encoder_length_loss))
            # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        # elif self.minimize_len and word_length_tensor is None:
        elif self.minimize_len:
            loss_fct = nn.L1Loss()
            real_length_loss = loss_fct(
                outputs.encoder_length_loss[0],
                torch.zeros_like(outputs.encoder_length_loss[0]),
            )

        loss = ((masked_lm_loss if masked_lm_loss is not None else 0.)
              + ((self.weight_len if self.weight_len is not None else 0.
                 ) * (real_length_loss if real_length_loss is not None else 0.))
        ) if ((masked_lm_loss is not None) or (real_length_loss is not None)) else None

        if not return_dict:
            output = (lm_logits,) + outputs[1:] + (masked_lm_loss, real_length_loss,)
            return (
                ((loss,) + output)
                if loss is not None else
                output)

        return AugSeq2SeqLMOutput(
            loss=loss,
            logits=lm_logits,
            past_key_values=outputs.past_key_values,
            decoder_hidden_states=outputs.decoder_hidden_states,
            decoder_attentions=outputs.decoder_attentions,
            cross_attentions=outputs.cross_attentions,
            encoder_last_hidden_state=outputs.encoder_last_hidden_state,
            encoder_last_hidden_out_attention_mask=outputs.encoder_last_hidden_out_attention_mask,
            encoder_length_loss=outputs.encoder_length_loss,
            encoder_pred_word_lengths=outputs.encoder_pred_word_lengths,
            encoder_hidden_states=outputs.encoder_hidden_states,
            encoder_attentions=outputs.encoder_attentions,
            masked_lm_loss=masked_lm_loss,
            real_length_loss=real_length_loss,
        )

    def prepare_inputs_for_generation(
        self,
        decoder_input_ids,
        word_length_tensor=None,
        alpha_values=None,
        past=None,
        attention_mask=None,
        head_mask=None,
        decoder_head_mask=None,
        cross_attn_head_mask=None,
        use_cache=None,
        encoder_outputs=None,
        **kwargs
    ):
        output_dict = super().prepare_inputs_for_generation(
            decoder_input_ids,
            past,
            attention_mask,
            head_mask,
            decoder_head_mask,
            cross_attn_head_mask,
            use_cache,
            encoder_outputs,
            **kwargs
        )

        return {
            "word_length_tensor": word_length_tensor,
            "alpha_values": alpha_values,
            **output_dict,
        }

    @staticmethod
    def _expand_inputs_for_generation(
        input_ids: torch.LongTensor,
        expand_size: int = 1,
        is_encoder_decoder: bool = False,
        attention_mask=None,
        encoder_outputs=None,
        **model_kwargs,
    ):
        expanded_return_idx = (
            torch.arange(input_ids.shape[0])
                 .view(-1, 1)
                 .repeat(1, expand_size)
                 .view(-1)
                 .to(input_ids.device)
        )
        input_ids, model_kwargs = (
            GenerationMixin._expand_inputs_for_generation(
                input_ids,
                expand_size,
                is_encoder_decoder,
                attention_mask,
                encoder_outputs,
                **model_kwargs,
        ))
        if "word_length_tensor" in model_kwargs and model_kwargs.get("word_length_tensor") is not None:
            word_length_tensor = model_kwargs["word_length_tensor"]
            model_kwargs["word_length_tensor"] = word_length_tensor.index_select(0, expanded_return_idx)
        elif "word_length_tensor" in model_kwargs and model_kwargs.get("word_length_tensor") is None:
            model_kwargs.pop("word_length_tensor")
        # TODO: maybe check model arch?
        if "alpha_values" in model_kwargs and model_kwargs.get("alpha_values") is not None:
            alpha_values = model_kwargs["alpha_values"]
            model_kwargs["alpha_values"] = alpha_values.index_select(0, expanded_return_idx)
        elif "alpha_values" in model_kwargs and model_kwargs.get("alpha_values") is None:
            model_kwargs.pop("alpha_values")
        # TODO: maybe check model arch?

        if "encoder_outputs" in model_kwargs:
            if "out_attention_mask" in encoder_outputs:
                out_attention_mask = model_kwargs["encoder_outputs"]["out_attention_mask"]
                model_kwargs["encoder_outputs"]["out_attention_mask"] = out_attention_mask.index_select(0, expanded_return_idx)

        return input_ids, model_kwargs
    # endregion <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<

class PreSentBartEncoder(BartEncoder):
    def __init__(self, config: BartConfig, embed_tokens: Optional[nn.Embedding] = None):
        super().__init__(config, embed_tokens)
        embed_dim = config.d_model
        self.collapse_n = getattr(config, "collapse_n", 0)
        self.config.collapse_n = self.collapse_n
        if self.collapse_n == -1:
            self.skip_cif = True
        else:
            self.skip_cif = False
        self.use_self = getattr(config, "use_self", False)
        self.config.use_self = self.use_self

        self.post_initialization(embed_dim, word_extractor=cif_function)

    def post_initialization(self, embed_dim, word_extractor=cif_function):
        # == encoder == #
        self.alpha_predictor = nn.Linear(embed_dim, 1)  # TODO: check!
        self.word_extractor = word_extractor
        self.length_predictor = nn.Linear(embed_dim, 1)

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        word_length_tensor: Optional[torch.LongTensor] = None,
        alpha_values: Optional[torch.FloatTensor] = None,
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
      ) -> Union[Tuple, AugBaseModelOutput]:

        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        # retrieve input_ids and inputs_embeds
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
            input_ids = input_ids.view(-1, input_shape[-1])
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale

        embed_pos = self.embed_positions(input_shape)

        hidden_states = inputs_embeds + embed_pos
        hidden_states = self.layernorm_embedding(hidden_states)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)

        # @@@@@@@@@@@@@@@@@@
        # # expand attention_mask
        # if attention_mask is not None:
        #     # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
        #     attention_mask = _expand_mask(attention_mask, inputs_embeds.dtype)
        # @@@@@@@@@@@@@@@@@@

        encoder_states = () if output_hidden_states else None
        all_attentions = () if output_attentions else None

        # check if head_mask has a correct number of layers specified if desired
        if head_mask is not None:
            if head_mask.size()[0] != (len(self.layers)):
                raise ValueError(
                    f"The head_mask should be specified for {len(self.layers)} layers, but it is for {head_mask.size()[0]}."
                )

        # ############################ #
        # >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>
        if self.skip_cif:
            hidden_states = hidden_states
            out_attention_mask = attention_mask
            length_loss = (
                torch.zeros(len(hidden_states)).float(),
                torch.zeros(len(hidden_states)).long(),)
            pred_word_lengths = None
        else:
            if self.collapse_n == 1:    # 1 for all
                word_length_tensor = torch.ones_like(word_length_tensor)
            elif self.collapse_n == 2:  # all for all
                word_length_tensor = (attention_mask.sum(-1).long()
                                      .to(word_length_tensor.device))
                                        # else: N for all
            (hidden_states, out_attention_mask, length_loss,
                          # out_attention_mask := (
                          #     encoder_output_attention_mask)
             pred_word_lengths,
             _, __) = self.sent_retriever(
                hidden_states,
                word_length_tensor=word_length_tensor,
                alpha_values=alpha_values,
                padding_mask=1 - attention_mask if attention_mask is not None else None)
        # <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
        # expand attention_mask
        if attention_mask is not None:
            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
            attention_mask = _expand_mask(out_attention_mask, inputs_embeds.dtype)
        # ############################ #

        for idx, encoder_layer in enumerate(self.layers):
            if output_hidden_states:
                encoder_states = encoder_states + (hidden_states,)
            # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
            dropout_probability = random.uniform(0, 1)
            if self.training and (dropout_probability < self.layerdrop):  # skip the layer
                layer_outputs = (None, None)
            else:
                if self.gradient_checkpointing and self.training:

                    def create_custom_forward(module):
                        def custom_forward(*inputs):
                            return module(*inputs, output_attentions)

                        return custom_forward

                    layer_outputs = torch.utils.checkpoint.checkpoint(
                        create_custom_forward(encoder_layer),
                        hidden_states,
                        attention_mask,
                        (head_mask[idx] if head_mask is not None else None),
                    )
                else:
                    layer_outputs = encoder_layer(
                        hidden_states,
                        attention_mask,
                        layer_head_mask=(head_mask[idx] if head_mask is not None else None),
                        output_attentions=output_attentions,
                    )

                hidden_states = layer_outputs[0]

            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        if output_hidden_states:
            encoder_states = encoder_states + (hidden_states,)

        if not return_dict:
            encoder_outputs = tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)
        else:
            encoder_outputs = BaseModelOutput(
                last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions
            )

        if not return_dict:
            return tuple(v for v in [
                hidden_states,
                out_attention_mask,
                encoder_outputs.hidden_states,
                length_loss,
                pred_word_lengths,
                encoder_outputs.attentions,
            # ] if v is not None)
            ] if v is not None)
        return AugBaseModelOutput(
            last_hidden_state=hidden_states,
            out_attention_mask=out_attention_mask,
            length_loss=length_loss,
            pred_word_lengths=pred_word_lengths,
            hidden_states=encoder_outputs.hidden_states,
            # # attentions=encoder_outputs.attentions,
            # attentions=encoder_outputs.attentions,
        )

    def sent_retriever(self,
        encoder__last_hidden_state,
        word_length_tensor=None,
        alpha_values=None,
        padding_mask=None,
        return_all=False,
        return_original=False,
      ):
        if alpha_values is None:
            alpha_values = self.alpha_predictor(encoder__last_hidden_state)
            alpha_values = alpha_values.squeeze(-1).sigmoid()  # B x S

        if word_length_tensor is None or self.use_self:
            # print("No given! self predict")
            word_length_tensor = alpha_values.sum(-1).long()
        else:
            # print("Wordlen given")
            # predicted_word_length_tensor = alpha_values.sum(-1).long()
            pass

        encoder__word_representations_CIF = (
            self.word_extractor(
                encoder__last_hidden_state,
                alpha=alpha_values,
                padding_mask=padding_mask,
                target_lengths=word_length_tensor,
            )
        )
        # TODO: Keep all CIF
        [encoder_word_representation] = encoder__word_representations_CIF['cif_out']
        [pred_word_lengths] = encoder__word_representations_CIF['alpha_sum']
        encoder_word_representation = encoder_word_representation.contiguous()
        # pred_word_lengths = pred_word_lengths.contiguous()
        # length_loss = 0.
        length_loss = ((pred_word_lengths, word_length_tensor,)
            if word_length_tensor is not None else
            (pred_word_lengths, pred_word_lengths,))
            # aliased as `encoder_word_representation`
            # FIXME: distributed problem!
            # TODO: add other CIF ouptuts!
        # length_loss = length_loss.contiguous()

        encoder_output_attention_mask = (
            # mask_generator(word_length_tensor)
            mask_generator(word_length_tensor)
            if word_length_tensor is not None else
            mask_generator(pred_word_lengths))
            # TODO: check length prediction!

        return (
            encoder_word_representation,
            encoder_output_attention_mask,
            length_loss,
            pred_word_lengths,
            encoder__word_representations_CIF if return_all else None,
            # encoder__last_hidden_state if return_original else None,
            encoder__last_hidden_state if return_original else None,
        )

class PreSentBart(SentBart):
    def __init__(self, config: BartConfig):
        super(SentBart, self).__init__(config)

        padding_idx, vocab_size = config.pad_token_id, config.vocab_size
        self.shared = nn.Embedding(vocab_size, config.d_model, padding_idx)

        self.encoder = PreSentBartEncoder(config, self.shared)
        self.decoder = BartDecoder(config, self.shared)

        # Initialize weights and apply final processing
        self.post_init()

class PreSentBartForConditionalGeneration(SentBartForConditionalGeneration):
    def __init__(self, config: BartConfig):
        super(SentBartForConditionalGeneration, self).__init__(config)
        self.model = PreSentBart(config)
        self.register_buffer(
            "final_logits_bias",
            torch.zeros((1, self.model.shared.num_embeddings)))
        self.lm_head = nn.Linear(
            config.d_model,
            self.model.shared.num_embeddings,
            bias=False)

        self.weight_len = getattr(config, "weight_len", None)
        self.config.weight_len = self.weight_len
        self.minimize_len = getattr(config, "minimize_len", False)
        self.config.minimize_len = self.minimize_len

        # Initialize weights and apply final processing
        self.post_init()


import pandas
from geopy.geocoders import ArcGIS

def geocode_addresses(df, header):
    arc = ArcGIS()
    def get_coords(row):
        return arc.geocode(row[header])[1]

    df[["Latitude", "Longitude"]] = df.apply(get_coords, axis=1).tolist()
    df["Latitude"] = df["Latitude"].round(decimals=6)
    df["Longitude"] = df["Longitude"].round(decimals=6)
    return df


# -*- coding: utf-8 -*-

import pygame
import pytmx
import settings as sett

class TiledMap:
    def __init__(self, filename):
        # pixelalpha=True to make sure the transparency go on
        tm = pytmx.load_pygame(filename, pixelalpha=True)

        # Width in tiles
        self.width = tm.width * tm.tilewidth
        self.height = tm.height * tm.tileheight

        # We will store all data in this variable
        self.tmxdata = tm

    def render(self, surface):
        # ti will be alias to this command
        ti = self.tmxdata.get_tile_image_by_gid

        my_gid = self.tmxdata.get_layer_by_name("Image").gid
        my_tile = ti(my_gid)
        surface.blit(my_tile, (0, 0))

        for layer in self.tmxdata.visible_layers:
            if isinstance(layer, pytmx.TiledTileLayer):
                for x, y, gid, in layer:
                    tile = ti(gid)
                    if tile:
                        surface.blit(tile, (x * self.tmxdata.tilewidth, y * self.tmxdata.tileheight))

    def make_map(self):
        temp_surface = pygame.Surface([self.width, self.height])
        self.render(temp_surface)
        return temp_surface

class Camera:
    def __init__(self, width, height):
        # As we can see camera is just a Rectangle
        self.camera = pygame.Rect(0, 0, width, height)
        self.width = width
        self.height = height

    # Apply offset to the entity
    def apply(self, entity):
        # move() gives new rectangle shifted by amount of an argument
        return entity.rect.move(self.camera.topleft)

    # Apply offset to the rectangle
    def apply_rect(self, rect):
        return rect.move(self.camera.topleft)

    def update(self, target):
        # We wanted centered on the screen
        x = -target.rect.x + (sett.WIDTH // 2)
        y = -target.rect.y + (sett.HEIGHT // 2)

        # Limit scrolling to map size
        x = min(0, x) # left
        y = min(0, y) # top
        x = max(-(self.width - sett.WIDTH), x) # right
        y = max(-(self.height - sett.HEIGHT), y) # bottom
        self.camera = pygame.Rect(x, y, self.width, self.height)


import logging
from typing import List

LOG = logging.getLogger(__name__)
SEX_MAP = {'male': 'M', 'female': 'F'}

class LimsHandler:

    @staticmethod
    def to_lims(customer: str, samples: List[dict]) -> List[dict]:
        """Convert order input to lims interface input."""
        samples_lims = []
        for sample in samples:
            LOG.debug(f"{sample['name']}: prepare LIMS input")
            samples_lims.append({
                'name': sample['name'],
                'container': sample.get('container') or 'Tube',
                'container_name': sample.get('container_name'),
                'well_position': sample.get('well_position'),
                'index_sequence': sample.get('index_sequence'),
                'udfs': {
                    'application': sample['application'],
                    'capture_kit': sample.get('capture_kit'),
                    'comment': sample.get('comment'),
                    'concentration': sample.get('concentration'),
                    'concentration_weight': sample.get('concentration_weight'),
                    'customer': customer,
                    'data_analysis': sample.get('data_analysis'),
                    'elution_buffer': sample.get('elution_buffer'),
                    'extraction_method': sample.get('extraction_method'),
                    'family_name': sample.get('family_name') or 'NA',
                    'formalin_fixation_time': sample.get('formalin_fixation_time'),
                    'index': sample.get('index'),
                    'index_number': sample.get('index_number'),
                    'organism': sample.get('organism'),
                    'organism_other': sample.get('organism_other'),
                    'pool': sample.get('pool'),
                    'post_formalin_fixation_time': sample.get('post_formalin_fixation_time'),
                    'priority': sample.get('priority') or 'standard',
                    'quantity': sample.get('quantity'),
                    'reference_genome': sample.get('reference_genome'),
                    'require_qcok': sample.get('require_qcok') or False,
                    'rml_plate_name': sample.get('rml_plate_name'),
                    'sex': SEX_MAP.get(sample.get('sex'), 'unknown'),
                    'source': sample.get('source') or 'NA',
                    'tissue_block_size': sample.get('tissue_block_size'),
                    'tumour': sample.get('tumour') or False,
                    'volume': sample.get('volume'),
                    'well_position_rml': sample.get('well_position_rml'),
                    'verified_organism': sample.get('verified_organism'),
                }
            })
        return samples_lims

    def process_lims(self, data: dict, samples: List[dict]):
        """Process samples to add them to LIMS."""
        samples_lims = self.to_lims(data['customer'], samples)
        project_name = data['ticket'] or data['name']
        project_data = self.lims.submit_project(project_name, samples_lims)
        lims_map = self.lims.get_samples(projectlimsid=project_data['id'], map_ids=True)
        return project_data, lims_map