import time

import torch
from peft import PeftModel, PeftConfig
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM

import gradio as gr
import speech_recognition as sr
from math import log2, pow
import os

#from scipy.fftpack import fft
import gc

peft_model_id="hackathon-somos-nlp-2023/SalpiBloomZ_15949_input_512-1b7"

config = PeftConfig.from_pretrained(peft_model_id)
model2 = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, 
                                               load_in_8bit=True, 
                                               # load_in_8bit_fp32_cpu_offload=True,
                                               device_map='auto')
tokenizer2 = AutoTokenizer.from_pretrained(peft_model_id)

model2 = PeftModel.from_pretrained(model2, peft_model_id)

Problema_tarjetaCredito= os.path.abspath("Problema_tarjetaCredito.ogg")
list_audios= [[Problema_tarjetaCredito]]

def gen_conversation(text,max_new_tokens=100): 
  text = "<SN>instruction: " + text + "\n "
  batch = tokenizer2(text, return_tensors='pt')
  batch = {k: v.to('cuda') for k, v in batch.items()}
  with torch.cuda.amp.autocast():
      output_tokens = model2.generate(**batch, 
                                    max_new_tokens=max_new_tokens, 
                                    eos_token_id= tokenizer2.eos_token_id,
                                    pad_token_id= tokenizer2.pad_token_id,
                                    bos_token_id= tokenizer2.bos_token_id,
                                    early_stopping = True, 
                                    no_repeat_ngram_size=2,
                                    repetition_penalty=1.2,
                                    temperature=.69,
                                    num_beams=3
                                    )
  gc.collect()
  torch.cuda.memory_summary(device=None, abbreviated=False)
  return tokenizer2.decode(output_tokens[0], skip_special_tokens=True).split("\n")[-1].replace("output:","")

conversacion = ""
description="""
Este proyecto nace de la idea de construir agentes virtuales especializados en temas académicos o nichos de negocio. 
Para esto fueron creados dos dataset principales, uno basado en habilidades blandas y otro basado en conversaciones empresariales donde un agente humano atiende a un cliente de la empresa en cuestión. 
Para ejecutar esta tarea, se han empleado múltiples LLMs, entre ellos:

- Alpaca native: https://huggingface.co/hackathon-somos-nlp-2023/salsapaca-native
- Bertin: https://huggingface.co/hackathon-somos-nlp-2023/bertin-gpt-j-6B-es-finetuned-salpaca
- Bloom: https://huggingface.co/hackathon-somos-nlp-2023/SalpiBloom-1b1
- Bloomz: https://huggingface.co/hackathon-somos-nlp-2023/SalpiBloomZ_15949_input_512-1b7
- Dolly: https://huggingface.co/hackathon-somos-nlp-2023/dolly-ArepacaV1
- Cerebras: https://huggingface.co/hackathon-somos-nlp-2023/Cerebras-Cabritax-softskills-espaGPT-1.3B

Teniendo como mejor resultado el modelo basado en BloomZ, el cual es el prensentado en esta demo.

La demo cuenta con dos tipos de inputs, mediante audio y mediante texto. 
Esto se hace con la idea de reflejar el hecho de que estos agentes virtuales pueden funcionar tanto para chatbots como para voicebots.
Con esto en mano, la empresa o institución que tenga acceso a esta tecnología, podría obtener un agente virtual con el conocimiento necesario para atender una conversación con el usuario final.
"""
def speech_to_text(audio_file, texto_adicional):
    global conversacion
    if audio_file is not None:
        # Lógica para entrada de audio
        r = sr.Recognizer()
        audio_data = sr.AudioFile(audio_file)
        with audio_data as source:
            audio = r.record(source)
        text_enrada=""
        
        texto_generado = r.recognize_google(audio, language="es-ES")
        texto_generado= f"[|Audio a texto|]:{texto_generado}\n" + "<br>[AGENTE]:"+gen_conversation(texto_generado,max_new_tokens=500)
        texto_generado = "<div style='color: #66b3ff;'>" + texto_generado + "</div><br>"
    else:
        texto_generado= f"[|Solo texto|]:{texto_adicional}\n" + "<br>[AGENTE]:"+gen_conversation(texto_adicional,max_new_tokens=500)
        texto_generado = "<div style='color: #66b3ff;'> " + texto_generado + "</div><br>"
    conversacion += texto_generado
    return conversacion

iface = gr.Interface(
    fn=speech_to_text,
    inputs=[gr.inputs.Audio(label="Voz", type="filepath"), gr.inputs.Textbox(label="Texto adicional")],
    outputs=gr.outputs.HTML(label=["chatbot","state"]),
    title="Chat bot para empresas.",
    description=description, #"Este modelo convierte la entrada de voz o texto y hace inferencia",
    examples=list_audios,
    theme="default",
    layout="vertical",
    allow_flagging=False,
    flagging_dir=None,
    server_name=None,
    server_port=None,
    live=False,
    capture_session=False
)

iface.launch()