import gradio as gr

import torch

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    pipeline
)

model_id = "somosnlp/gemma-7b-it-legal-refugee-v0.1.1"
tokenizer_id = "philschmid/gemma-tokenizer-chatml"

tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
# Cargamos el modelo en 4 bits para agilizar la inferencia
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16,
    bnb_4bit_use_double_quant=True,
)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="auto",
    quantization_config=quantization_config,
)

print(f"Model loaded in {model.device}")
# Generamos el pipeline de generación de texto
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
# Definimos el eos token
eos_token = tokenizer("<|im_end|>",add_special_tokens=False)["input_ids"][0]

def generate_inference(prompt):
    prompt = pipe.tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True)
    outputs = pipe(prompt, max_new_tokens=128, do_sample=True, num_beams=1, temperature=0.7, top_k=50, top_p=0.95, 
                   max_time= 180, eos_token_id=eos_token)
    return outputs[0]['generated_text'][len(prompt):].strip()

demo = gr.Interface(
    generate_inference,
    inputs=[
        gr.Textbox(
            label="🌐 Pregunta a resolver",
            info="Introduce la duda o pregunta que la que deseas obtener información.",
            value="Pregunta de prueba",
            interactive=True,
        )
    ],
    outputs=[
        gr.Textbox(
            label="📰 Respuesta e información",
            interactive=False,
            placeholder="Aquí aparecerá la respuesta o información solicitada",
        )
    ],
    title="QUESTION ANSWERING LEGAL REFUGIADOS v1.0",
    theme="JohnSmith9982/small_and_pretty",
    description="Apoyo y consulta acerca de leyes y protección a asiliados y refugiados",
    #article=article,
    #cache_examples=False,
    concurrency_limit=1,
    examples=[
        "¿Quién es un solicitante de asilo/ protección internacional?",
        "¿Cuáles son los derechos y obligaciones de un solicitante de protección internacional?",
        "¿Dónde se puede solicitar asilo/ protección internacional?",
        "una vez declarada la protección temporal por el Consejo de la Unión Europea, que responsabilidades del Gobierno español?",
    ],
    submit_btn="Generar respuesta",
    stop_btn="Detener respuesta",
    clear_btn="Limpiar",
    allow_flagging="never",
)

#demo.queue(max_size=None)
demo.launch(debug=True, share=False)