<jupyter_start><jupyter_text><jupyter_code>import os

os.environ["CUDA_VISIBLE_DEVICES"] = '0'
os.environ["WANDB_DISABLED"] = 'True'
 
os.environ['CUDA_LAUNCH_BLOCKING'] = "1"
!pip install collections
!pip install datasets
!pip install transformers[sentencepiece]
!pip install tensorboard
!pip install torch

import torch
from tqdm.auto import tqdm
import numpy as np
import pandas as pd
from datasets import Dataset
from transformers import AutoModelForQuestionAnswering, TrainingArguments, Trainer, AutoTokenizer, default_data_collator
import collections
import datetime
import sys
from google.colab import drive

from torch.utils.tensorboard import SummaryWriter
writer = SummaryWriter()

drive.mount('/content/drive')

BASE_PATH = "/content/drive/MyDrive/"

os.getcwd()


torch.cuda.is_available()
torch.cuda.device_count()

"""import random
seed_num = 123
random.seed(seed_num)
torch.manual_seed(seed_num)
np.random.seed(seed_num)"""
max_length = 256 # The maximum length of a feature (question and context)
doc_stride = 64# The authorized overlap between two part of the context when splitting it is needed.
batch_size = 4
from transformers import AutoModel, AutoTokenizer
hindi_tokenizer = AutoTokenizer.from_pretrained(BASE_PATH + "roberta-large-wechsel-hindi")
hindi_model = AutoModelForQuestionAnswering.from_pretrained(BASE_PATH + "roberta-large-wechsel-hindi")

"""tamil_tokenizer = AutoTokenizer.from_pretrained(BASE_PATH + "roberta-large-wechsel-tamil")
tamil_model = AutoModelForQuestionAnswering.from_pretrained(BASE_PATH + "roberta-large-wechsel-tamil")

pad_on_right = tamil_tokenizer.padding_side == "right"""

pad_on_right = hindi_tokenizer.padding_side == "right"


hindi_model.roberta.config.type_vocab_size = 2
single_emb = hindi_model.roberta.embeddings.token_type_embeddings
hindi_model.roberta.embeddings.token_type_embeddings = torch.nn.Embedding(2, single_emb.embedding_dim)
hindi_model.roberta.embeddings.token_type_embeddings.weight = torch.nn.Parameter(single_emb.weight.repeat([2, 1]))

"""tamil_model.roberta.config.type_vocab_size = 2
single_emb = tamil_model.roberta.embeddings.token_type_embeddings
tamil_model.roberta.embeddings.token_type_embeddings = torch.nn.Embedding(2, single_emb.embedding_dim)
tamil_model.roberta.embeddings.token_type_embeddings.weight = torch.nn.Parameter(single_emb.weight.repeat([2, 1]))"""
df_train_base = pd.read_csv(BASE_PATH + "train.csv")
df_test = pd.read_csv(BASE_PATH + "test.csv")
df_sub = pd.read_csv(BASE_PATH + "sample_submission.csv")
df_train_base = pd.concat([df_train_base], axis=0, ignore_index=True)
def prepare_train_features_hindi(examples):

    examples["question"] = [q.lstrip() for q in examples["question"]]
    
    tokenized_examples = hindi_tokenizer(
        examples["question" if pad_on_right else "context"],
        examples["context" if pad_on_right else "question"],
        truncation="only_second" if pad_on_right else "only_first",
        max_length=max_length,
        stride=doc_stride,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length",
    )

    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
    offset_mapping = tokenized_examples.pop("offset_mapping")

    tokenized_examples["start_positions"] = []
    tokenized_examples["end_positions"] = []

    for i, offsets in enumerate(offset_mapping):
        input_ids = tokenized_examples["input_ids"][i]
        cls_index = input_ids.index(hindi_tokenizer.cls_token_id)
        sequence_ids = tokenized_examples.sequence_ids(i)
        sample_index = sample_mapping[i]
        answers = examples["answers"][sample_index]
        
        if len(answers["answer_start"]) == 0:
            tokenized_examples["start_positions"].append(cls_index)
            tokenized_examples["end_positions"].append(cls_index)
        else:
            
            start_char = answers["answer_start"][0]
            end_char = start_char + len(answers["text"][0])
            
            token_start_index = 0
            while sequence_ids[token_start_index] != (1 if pad_on_right else 0):
                token_start_index += 1

            token_end_index = len(input_ids) - 1
            while sequence_ids[token_end_index] != (1 if pad_on_right else 0):
                token_end_index -= 1

            if not (offsets[token_start_index][0] <= start_char and offsets[token_end_index][1] >= end_char):
                tokenized_examples["start_positions"].append(cls_index)
                tokenized_examples["end_positions"].append(cls_index)
            else:
                while token_start_index < len(offsets) and offsets[token_start_index][0] <= start_char:
                    token_start_index += 1
                tokenized_examples["start_positions"].append(token_start_index - 1)
                while offsets[token_end_index][1] >= end_char:
                    token_end_index -= 1
                tokenized_examples["end_positions"].append(token_end_index + 1)

    return tokenized_examples

def prepare_train_features_tamil(examples):

    examples["question"] = [q.lstrip() for q in examples["question"]]
    
    tokenized_examples = tamil_tokenizer(
        examples["question" if pad_on_right else "context"],
        examples["context" if pad_on_right else "question"],
        truncation="only_second" if pad_on_right else "only_first",
        max_length=max_length,
        stride=doc_stride,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length",
    )

    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
    offset_mapping = tokenized_examples.pop("offset_mapping")

    tokenized_examples["start_positions"] = []
    tokenized_examples["end_positions"] = []

    for i, offsets in enumerate(offset_mapping):
        input_ids = tokenized_examples["input_ids"][i]
        cls_index = input_ids.index(tamil_tokenizer.cls_token_id)
        sequence_ids = tokenized_examples.sequence_ids(i)
        sample_index = sample_mapping[i]
        answers = examples["answers"][sample_index]
        
        if len(answers["answer_start"]) == 0:
            tokenized_examples["start_positions"].append(cls_index)
            tokenized_examples["end_positions"].append(cls_index)
        else:
            
            start_char = answers["answer_start"][0]
            end_char = start_char + len(answers["text"][0])
            
            token_start_index = 0
            while sequence_ids[token_start_index] != (1 if pad_on_right else 0):
                token_start_index += 1

            token_end_index = len(input_ids) - 1
            while sequence_ids[token_end_index] != (1 if pad_on_right else 0):
                token_end_index -= 1

            if not (offsets[token_start_index][0] <= start_char and offsets[token_end_index][1] >= end_char):
                tokenized_examples["start_positions"].append(cls_index)
                tokenized_examples["end_positions"].append(cls_index)
            else:
                while token_start_index < len(offsets) and offsets[token_start_index][0] <= start_char:
                    token_start_index += 1
                tokenized_examples["start_positions"].append(token_start_index - 1)
                while offsets[token_end_index][1] >= end_char:
                    token_end_index -= 1
                tokenized_examples["end_positions"].append(token_end_index + 1)

    return tokenized_examples

def prepare_validation_features(examples):
    examples["question"] = [q.lstrip() for q in examples["question"]]
    tokenized_examples = tokenizer(
        examples["question" if pad_on_right else "context"],
        examples["context" if pad_on_right else "question"],
        truncation="only_second" if pad_on_right else "only_first",
        max_length=max_length,
        stride=doc_stride,
        return_overflowing_tokens=True,
        return_offsets_mapping=True,
        padding="max_length",
    )
    
    sample_mapping = tokenized_examples.pop("overflow_to_sample_mapping")
    tokenized_examples["example_id"] = []

    for i in range(len(tokenized_examples["input_ids"])):
        sequence_ids = tokenized_examples.sequence_ids(i)
        context_index = 1 if pad_on_right else 0
        sample_index = sample_mapping[i]
        tokenized_examples["example_id"].append(examples["id"][sample_index])
        tokenized_examples["offset_mapping"][i] = [
            (o if sequence_ids[k] == context_index else None)
            for k, o in enumerate(tokenized_examples["offset_mapping"][i])
        ]

    return tokenized_examples


def postprocess_qa_predictions(examples, features, raw_predictions, n_best_size = 20, max_answer_length = 30):
    all_start_logits, all_end_logits = raw_predictions
    # Build a map example to its corresponding features.
    example_id_to_index = {k: i for i, k in enumerate(examples["id"])}
    features_per_example = collections.defaultdict(list)
    for i, feature in enumerate(features):
        features_per_example[example_id_to_index[feature["example_id"]]].append(i)

    # The dictionaries we have to fill.
    predictions = collections.OrderedDict()

    # Logging.
    print(f"Post-processing {len(examples)} example predictions split into {len(features)} features.")

    # Let's loop over all the examples!
    for example_index, example in enumerate(tqdm(examples)):
        # Those are the indices of the features associated to the current example.
        feature_indices = features_per_example[example_index]

        min_null_score = None # Only used if squad_v2 is True.
        valid_answers = []
        
        context = example["context"]
        # Looping through all the features associated to the current example.
        for feature_index in feature_indices:
            # We grab the predictions of the model for this feature.
            start_logits = all_start_logits[feature_index]
            end_logits = all_end_logits[feature_index]
            # This is what will allow us to map some the positions in our logits to span of texts in the original
            # context.
            offset_mapping = features[feature_index]["offset_mapping"]

            # Update minimum null prediction.
            cls_index = features[feature_index]["input_ids"].index(tokenizer.cls_token_id)
            feature_null_score = start_logits[cls_index] + end_logits[cls_index]
            if min_null_score is None or min_null_score < feature_null_score:
                min_null_score = feature_null_score

            # Go through all possibilities for the `n_best_size` greater start and end logits.
            start_indexes = np.argsort(start_logits)[-1 : -n_best_size - 1 : -1].tolist()
            end_indexes = np.argsort(end_logits)[-1 : -n_best_size - 1 : -1].tolist()
            for start_index in start_indexes:
                for end_index in end_indexes:
                    # Don't consider out-of-scope answers, either because the indices are out of bounds or correspond
                    # to part of the input_ids that are not in the context.
                    if (
                        start_index >= len(offset_mapping)
                        or end_index >= len(offset_mapping)
                        or offset_mapping[start_index] is None
                        or offset_mapping[end_index] is None
                    ):
                        continue
                    # Don't consider answers with a length that is either < 0 or > max_answer_length.
                    if end_index < start_index or end_index - start_index + 1 > max_answer_length:
                        continue

                    start_char = offset_mapping[start_index][0]
                    end_char = offset_mapping[end_index][1]
                    valid_answers.append(
                        {
                            "score": start_logits[start_index] + end_logits[end_index],
                            "text": context[start_char: end_char]
                        }
                    )
        
        if len(valid_answers) > 0:
            best_answer = sorted(valid_answers, key=lambda x: x["score"], reverse=True)[0]
        else:
            # In the very rare edge case we have not a single non-null prediction, we create a fake prediction to avoid
            # failure.
            best_answer = {"text": "", "score": 0.0}
        
        # Let's pick our final answer: the best one or the null answer (only for squad_v2)
        predictions[example["id"]] = best_answer["text"]

    return predictions

def jaccard(row): 
    a = set(row['answer'].lower().split()) 
    b = set(row['prediction'].lower().split())
    c = a.intersection(b)
    return float(len(c)) / (len(a) + len(b) - len(c))

def convert_answers(row):
    return {'answer_start': [row['answer_start']], 'text': [row['answer_text']]}
df_train_base['answers'] = df_train_base[['answer_start', 'answer_text']].apply(convert_answers, axis=1)
df_train_base = df_train_base.sample(frac=1, random_state=114514).copy()

df_train_base_tamil = df_train_base.loc[(df_train_base['language'] == "tamil")]
df_train_base_hindi= df_train_base.loc[(df_train_base['language'] == "hindi")]

df_train_tamil = df_train_base_tamil.reset_index(drop=True)
df_valid_tamil = df_train_base_tamil[-64:].reset_index(drop=True)

df_train_hindi= df_train_base_hindi.reset_index(drop=True)
df_valid_hindi = df_train_base_hindi[-64:].reset_index(drop=True)

train_dataset_hindi = Dataset.from_pandas(df_train_hindi)
valid_dataset_hindi = Dataset.from_pandas(df_valid_hindi)

train_dataset_tamil = Dataset.from_pandas(df_train_tamil)
valid_dataset_tamil = Dataset.from_pandas(df_valid_tamil)

"""tokenized_train_tamil_ds = train_dataset_tamil.map(prepare_train_features_tamil, batched=True, remove_columns=train_dataset_tamil.column_names)
tokenized_valid_tamil_ds = valid_dataset_tamil.map(prepare_train_features_tamil, batched=True, remove_columns=train_dataset_tamil.column_names)"""

tokenized_train_hindi_ds = train_dataset_hindi.map(prepare_train_features_hindi, batched=True, remove_columns=train_dataset_hindi.column_names)
tokenized_valid_hindi_ds = valid_dataset_hindi.map(prepare_train_features_hindi, batched=True, remove_columns=train_dataset_hindi.column_names)
print(len(df_train_base_hindi))
print(len(train_dataset_hindi))
print(len(tokenized_train_hindi_ds))
def get_trainer_tamil():
    model = tamil_model
    args = TrainingArguments(
    f"chaii-qa/" + "tamil",
    evaluation_strategy = "epoch",
    learning_rate=2e-5,
    gradient_accumulation_steps=8 // batch_size,
    warmup_ratio=0.1,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    num_train_epochs=4, # Please set this as 2 for reproduction
    weight_decay=0.01,
    dataloader_num_workers=2,
    fp16=True,
    fp16_opt_level='O3',
    #fp16_backend='apex'
    )

    trainer = Trainer(model, args,
                      train_dataset= tokenized_train_tamil_ds,
                      eval_dataset= tokenized_valid_tamil_ds,
                      data_collator=default_data_collator,
                      tokenizer=tamil_tokenizer
    )
    
    return trainer
def get_trainer_hindi():
    model = hindi_model
    args = TrainingArguments(
    f"chaii-qa/" + "hindi",
    evaluation_strategy = "epoch",
    learning_rate=2e-5,
    warmup_ratio=0.1,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    save_strategy="steps",
    num_train_epochs=4, # Please set this as 2 for reproduction
    weight_decay=0.01,
    dataloader_num_workers=2,
    fp16=True,
    fp16_opt_level='O3',
    #fp16_backend='apex'
    )

    trainer = Trainer(model, args,
                      train_dataset= tokenized_train_hindi_ds,
                      eval_dataset= tokenized_valid_hindi_ds,
                      data_collator=default_data_collator,
                      tokenizer=hindi_tokenizer
    )
    
    return trainer


print(hindi_model.roberta.embeddings.word_embeddings)

print(len(hindi_tokenizer))
%load_ext tensorboard
%tensorboard --logdir /chaii-qa/hindi/
batch_size = 12
trainer_hindi = get_trainer_hindi()
"""trainer_tamil = get_trainer_tamil()"""



trainer_hindi.train()

trainer_hindi.save_model("hindi/")



"""trainer_tamil.train()
trainer_tamil.save_model("tamil/")"""


"""validation_features = valid_dataset.map(prepare_validation_features, batched=True,remove_columns=valid_dataset.column_names)

len(validation_features)

validation_features

valid_feats_small = validation_features.map(lambda example: example, remove_columns=['example_id', 'offset_mapping'])
valid_feats_small

raw_predictions = trainer_hindi.predict(valid_feats_small)

examples = valid_dataset
features = validation_features

example_id_to_index = {k: i for i, k in enumerate(examples["id"])}
features_per_example = collections.defaultdict(list)
for i, feature in enumerate(features):
    features_per_example[example_id_to_index[feature["example_id"]]].append(i)
    
final_predictions = postprocess_qa_predictions(valid_dataset, validation_features, raw_predictions.predictions)
references = [{"id": ex["id"], "answer": ex["answers"]['text'][0]} for ex in valid_dataset]


res = pd.DataFrame(references)
res['prediction'] = res['id'].apply(lambda r: final_predictions[r])
res['jaccard'] = res[['answer', 'prediction']].apply(jaccard, axis=1)
res


print(res['jaccard'].mean())"""<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09727937814357568, 'fraction_numerical': 0.007487425697302241, 'mean_word_length': 3.474936061381074, 'pattern_counts': {'":': 6, '<': 14, '<?xml version=': 0, '>': 14, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '136928', 'license': 'WTFPL', 'n_tokens_mistral': 5614, 'n_tokens_neox': 5372, 'n_words': 990, 'path': 'indicdataonly.ipynb', 'repo_name': 'adhi-thirumala/Science-Fair2022', 'size': '17496'}