edumunozsala's picture
Update app.py
dfe5f78
raw
history blame
2.71 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.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)