|
import gradio as gr |
|
import os |
|
import spaces |
|
from huggingface_hub import login |
|
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer |
|
from threading import Thread |
|
|
|
|
|
hf_token = os.getenv("LLAMA31") |
|
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.") |
|
|
|
|
|
bnb_config = BitsAndBytesConfig(load_in_4bit=True) |
|
|
|
|
|
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, |
|
device_map="auto" |
|
) |
|
|
|
|
|
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. |
|
"""} |
|
|
|
|
|
def chat_fn(multimodal_message): |
|
|
|
question = multimodal_message["text"] |
|
|
|
|
|
conversation = [{"role": "user", "content": question}] |
|
|
|
|
|
input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt") |
|
input_ids = input_ids.to(model.device) |
|
|
|
|
|
streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True) |
|
|
|
|
|
generate_kwargs = dict( |
|
input_ids=input_ids, |
|
streamer=streamer, |
|
max_new_tokens=500, |
|
do_sample=True, |
|
temperature=0.7, |
|
) |
|
|
|
|
|
t = Thread(target=model.generate, kwargs=generate_kwargs) |
|
t.start() |
|
|
|
|
|
message = "" |
|
for text in streamer: |
|
message += text |
|
yield message |
|
|
|
|
|
with gr.Blocks() as demo: |
|
|
|
gr.Markdown("# 🔍 Chatbot Analizador de Documentos") |
|
|
|
|
|
response = gr.Textbox(lines=5, label="Respuesta") |
|
|
|
|
|
chat = gr.MultimodalTextbox(file_types=["image"], interactive=True, |
|
show_label=False, placeholder="Sube una imagen del documento haciendo clic en '+' y haz una pregunta.") |
|
|
|
|
|
chat.submit(chat_fn, inputs=chat, outputs=response) |
|
|
|
|
|
if __name__ == "__main__": |
|
demo.launch() |
|
|