edumunozsala's picture
Update app.py
88cebd1
raw
history blame
2.93 kB
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.3, 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="¿Quién es un solicitante de asilo/ protección internacional?",
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="ParityError/LimeFace",
description="Apoyo y consulta acerca de leyes y procedimientos de protección a asiliados y refugiados",
article="Escribe tu duda o cuestión y trataremos de darte una respuesta y ayuda que te permita resolver tu cuestión",
thumbnail="markus-winkler-Je1MDuITTF4-unsplash.jpg",
#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)