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

#import streamlit as st   #to use the shared secrets
from huggingface_hub import login
import gradio as gr
#from transformers import pipeline 
import torch
from utils import *
from presets import *
from transformers import Trainer, TrainingArguments

#####################################################
#Hilfsfunktionen für das training
#####################################################
#Datensets in den Tokenizer schieben...
def tokenize_function(examples):
    return tokenizer(examples["text"])

    
#Funktion, die den gegebenen Text aus dem Datenset gruppiert
def group_texts(examples):
    # Concatenate all texts.
    concatenated_examples = {k: sum(examples[k], []) for k in examples.keys()}
    total_length = len(concatenated_examples[list(examples.keys())[0]])
    # We drop the small remainder, we could add padding if the model supported it instead of this drop, you can
    # customize this part to your needs.
    total_length = (total_length // block_size) * block_size
    # Split by chunks of max_len.
    result = {
        k: [t[i : i + block_size] for i in range(0, total_length, block_size)]
        for k, t in concatenated_examples.items()
    }
    result["labels"] = result["input_ids"].copy()
    return result



###################################################################################
#Access-Token (in Secrets)

#access_token = st.secrets["ACCESS_TOKEN"]
access_token = "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"
login(token="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")
#Modelle und Tokenizer

#Alternativ mit beliebigen Modellen:
base_model = "project-baize/baize-v2-7b"  #load_8bit = False (in load_tokenizer_and_model)
#base_model = "TheBloke/airoboros-13B-HF"  #load_8bit = False (in load_tokenizer_and_model)
#base_model = "EleutherAI/gpt-neo-1.3B"    #load_8bit = False (in load_tokenizer_and_model)
#base_model = "TheBloke/airoboros-13B-HF"   #load_8bit = True
tokenizer,model,device = load_tokenizer_and_model(base_model, access_token, False)

dataset_neu = daten_laden("alexkueck/tis") 
#dataset_neu = daten_laden("EleutherAI/pile")

#############################################
#Vorbereiten für das Training der neuen Daten
#############################################

#alles zusammen auf das neue datenset anwenden - batched = True und 4 Prozesse, um die Berechnung zu beschleunigen. Die "text" - Spalte braucht man anschließend nicht mehr, daher weglassen.
tokenized_datasets = dataset_neu.map(tokenize_function, batched=True, num_proc=4, remove_columns=["id","text"])

print (tokenized_datasets["train"][4])
#den Text nun zusammenführen (concatenieren) und anschließend in kleine Häppchen aufteilen (block_size=128), die verarbeitet werden können
#das macht die map-Funktion und das Attribut batched = True
#man könnte das weglassen, wenn jeder Satz einzeln gegeben wurde in den Texten...
#eigentlich nimmt man als block_size die max. Länge in der das Model trainiert wurde -> könnte aber zu groß sein für den RAm der GPU , daher hier 128 gewählt
# block_size = tokenizer.model_max_length
block_size = 128

#nochmal die map-Funktion auf das bereits tokenisierte Datenset anwenden
#die bereits tokenisierten Datensatze ändern sich dadurch: die samples enthalten nun Mengen aus block_size Tokens
lm_datasets = tokenized_datasets.map(
    group_texts,
    batched=True,
    batch_size=1000,
    num_proc=4,
)

#die Daten wurden nun "gereinigt" und für das Model vorbereitet.
#z.B. anschauen mit: tokenizer.decode(lm_datasets["train"][1]["input_ids"])

####################################################
#Training
####################################################

#Training Args
model_name = base_model.split("/")[-1]
training_args = TrainingArguments(
    "/model",
    evaluation_strategy = "epoch",
    learning_rate=2e-5,
    weight_decay=0.01,
    push_to_hub=True,
)

############################################
def trainieren_neu(name):
    #Trainer zusammenstellen
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=lm_datasets["train"],
    )

    #trainer ausführen
    trainer.train()

    #in den Hub laden
    trainer.push_to_hub()






#######################################################################
#Darstellung mit Gradio

with gr.Blocks() as demo:
    name = gr.Textbox(label="Model")
    output = gr.Textbox(label="Output Box")
    start_btn = gr.Button("Start")
    start_btn.click(fn=trainieren_neu, inputs=name, outputs=output, api_name="trainieren_neu")


demo.launch()