Spaces:
Runtime error
Runtime error
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", | |
) | |
], | |
headline="LEGAL REFUGIADOS 1.0 ", | |
#thumbnail="https://huggingface.co/datasets/Iker/NoticIA/resolve/main/assets/logo.png", | |
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=False, | |
) | |
#demo.queue(max_size=None) | |
demo.launch(debug=True, share=False) | |