from unsloth import FastLanguageModel

# 4bit pre quantized models we support for 4x faster downloading + no OOMs.
fourbit_models = [
    "unsloth/mistral-7b-instruct-v0.2-bnb-4bit",
    "unsloth/gemma-7b-it-bnb-4bit",
] # More models at https://huggingface.co/unsloth

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "unsloth/Meta-Llama-3.1-8B-Instruct",
    max_seq_length = 8192,
    load_in_4bit = True,
    # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
)
####
from transformers import TextStreamer
from unsloth.chat_templates import get_chat_template
tokenizer = get_chat_template(
    tokenizer,
    chat_template = "llama-3.1",
    mapping = {"role" : "from", "content" : "value", "user" : "human", "assistant" : "gpt"}, # ShareGPT style
)
FastLanguageModel.for_inference(model) # Enable native 2x faster inference
####
from datetime import datetime
import pendulum
from babel.dates import format_date

def get_day_month():
  # using São Paulo time zone
  SP = pendulum.timezone('America/Sao_Paulo')
  #
  date = datetime.now(SP).strftime('%Y/%m/%d')
  date = datetime.strptime(date, "%Y/%m/%d")
  # Spanish locale
  spanish_date = format_date(date, "dd 'de' MMMM", locale='es')

  return (spanish_date)

    #return f"""The {quantity} {animal}s from {" and ".join(countries)} went to the {place} where they {" and ".join(activity_list)} until the {"morning" if morning else "night"}"""
def sentence_builder(signo, estilo,tema):
  dia_texto = get_day_month()
  texto_request = None
  #if tema == "General":
  #  texto_request = f"""Actua como un experto en astrología. Hoy es {dia_texto} Soy del {signo}, haz una predicción sucinta, resumida, y precisa con un estilo {estilo}. No agregues ningun comentario en relación a la astrología"""
  #else:
  match tema:
    case "Salud":
      texto_request = f"""Actua como un experto en astrología. Hoy es {dia_texto} Soy del {signo}, haz una predicción sucinta, resumida, y precisa sobre {tema}, con un estilo {estilo}. No agregues ningun comentario en relación a la astrología"""
    case "Económico":
      texto_request = f"""Actua como un experto en astrología. Hoy es {dia_texto} Soy del {signo}, haz una predicción sucinta, resumida, y precisa sobre lo {tema}, con un estilo {estilo}. No agregues ningun comentario en relación a la astrología"""
    case "Emocional":
      texto_request = f"""Actua como un experto en astrología. Hoy es {dia_texto} Soy del {signo}, haz una predicción sucinta, resumida, y precisa sobre mi vida {tema}, con un estilo {estilo}. No agregues ningun comentario en relación a la astrología"""
    case "Amistades":
      texto_request = f"""Actua como un experto en astrología. Hoy es {dia_texto} Soy del {signo}, haz una predicción sucinta, resumida, y precisa sobre mis {tema}, con un estilo {estilo}. No agregues ningun comentario en relación a la astrología"""
    case "Familia":
      texto_request = f"""Actua como un experto en astrología. Hoy es {dia_texto} Soy del {signo}, haz una predicción sucinta, resumida, y precisa sobre mi vida {tema}r, con un estilo {estilo}. No agregues ningun comentario en relación a la astrología"""
    case _:
      texto_request = f"""Actua como un experto en astrología. Hoy es {dia_texto} Soy del {signo}, haz una predicción sucinta, resumida, y precisa con un estilo {estilo}. No agregues ningun comentario en relación a la astrología"""
  # "Emocional","Amistades","Familia"
  messages = [
    {"from": "human", "value": texto_request},
  ]
  inputs = tokenizer.apply_chat_template(messages, tokenize = True, add_generation_prompt = True, return_tensors = "pt").to("cuda")
  #
  model_output = model.generate(input_ids = inputs, streamer = text_streamer, max_new_tokens = 1024, use_cache = True)
  #
  model_output = tokenizer.decode(model_output[0].tolist(),skip_special_tokens=True)
  #print("===============")
  #print(model_output.split("assistant"))
  return model_output.split("assistant")[1]

###########
import gradio as gr


demo = gr.Interface(
    sentence_builder,
    [
        gr.Dropdown(
            ["Aries", "Tauro", "Géminis", "Cáncer", "Leo", "Virgo", "Libra", "Escorpio", "Sagitario", "Capricornio", "Acuario", "Piscis"], label="Signo", info="Escoge tu signo del Zodiaco"
        ),
        gr.Dropdown(
            ["Alegre","Serio","Sarcástico","Depresivo","Coqueto","Sensual"], label="Estilo", info="Elige el estilo que más gustes!",value="Alegre"
        ),
        gr.Dropdown(
            ["General","Salud","Económico","Emocional","Amistades","Familia"], label="Tema", info="Elige el aspecto de tu vida que te preocupa.",value="General"
        ),
    ],
    "textbox"
)

if __name__ == "__main__":
    demo.launch()