UPN / app.py
markIA23's picture
Update app.py
e6f4112 verified
raw
history blame
4.13 kB
import gradio as gr
import os
import spaces
from huggingface_hub import login
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
from threading import Thread
# Autenticación en Hugging Face
hf_token = os.getenv("LLAMA31") # Asegúrate de configurar la variable de entorno HF_TOKEN con tu token
if hf_token:
login(token=hf_token)
else:
raise ValueError("Hugging Face token no encontrado. Asegúrate de que la variable de entorno HF_TOKEN esté configurada.")
# Configuración para cargar el modelo en 4 bits utilizando bitsandbytes
bnb_config = BitsAndBytesConfig(load_in_4bit=True)
# Cargar el tokenizador y el modelo cuantizado desde Hugging Face
model_id = "meta-llama/Meta-Llama-3.1-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
quantization_config=bnb_config, # Cuantización en 4 bits
device_map="auto"
)
# Definimos un mensaje inicial que simula un rol del sistema y contiene instrucciones sobre cómo responder, adaptado al español
system_message = {
"role": "system",
"content": """
Eres un asistente útil.
Recibirás una pregunta y un conjunto de respuestas junto con una puntuación de confianza entre 0 y 1 para cada respuesta.
Tu trabajo es convertir esta información en una respuesta corta y coherente.
Por ejemplo:
Pregunta: "¿A quién se le está facturando?", respuesta: {"answer": "John Doe", "confidence": 0.98}
Deberías responder algo como:
Con un alto grado de confianza, puedo decir que se le está facturando a John Doe.
Pregunta: "¿Cuál es el total de la factura?", respuesta: [{"answer": "154.08", "confidence": 0.75}, {"answer": "155", "confidence": 0.25}]
Deberías responder algo como:
Creo que el total de la factura es de $154.08 aunque también podría ser $155.
"""}
# Definimos la función de inferencia utilizando el modelo Meta Llama 3.1 8B Instruct
def chat_fn(multimodal_message):
# Extraemos el texto de la pregunta del mensaje proporcionado por el usuario
question = multimodal_message["text"]
# Construimos el mensaje para el modelo
conversation = [{"role": "user", "content": question}]
# Generamos los IDs de entrada utilizando el tokenizador del modelo
input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt")
input_ids = input_ids.to(model.device)
# Configuramos el streamer para la generación progresiva de texto
streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
# Configuramos los argumentos de generación
generate_kwargs = dict(
input_ids=input_ids,
streamer=streamer,
max_new_tokens=500, # Ajusta esto según tus necesidades
do_sample=True,
temperature=0.7, # Ajusta la temperatura según tus necesidades
)
# Iniciamos la generación de texto en un hilo separado
t = Thread(target=model.generate, kwargs=generate_kwargs)
t.start()
# Iteramos sobre los tokens generados y construimos la respuesta
message = ""
for text in streamer:
message += text
yield message
# Usamos la clase 'Blocks' de Gradio para definir la interfaz de usuario
with gr.Blocks() as demo:
# Título de la aplicación en español
gr.Markdown("# 🔍 Chatbot Analizador de Documentos")
# Cuadro de texto para mostrar la respuesta generada, etiquetado en español
response = gr.Textbox(lines=5, label="Respuesta")
# Campo de texto multimodal para que el usuario suba un archivo e ingrese una pregunta, en español
chat = gr.MultimodalTextbox(file_types=["image"], interactive=True,
show_label=False, placeholder="Sube una imagen del documento haciendo clic en '+' y haz una pregunta.")
# Se asigna la función chat_fn para que se ejecute cuando el usuario envíe un mensaje en chat
chat.submit(chat_fn, inputs=chat, outputs=response)
# Lanza la aplicación si este archivo es ejecutado directamente
if __name__ == "__main__":
demo.launch()