Spaces:
Sleeping
Sleeping
File size: 3,196 Bytes
4ad7141 332e828 4ad7141 332e828 4ad7141 332e828 4ad7141 332e828 4ad7141 332e828 4ad7141 332e828 4ad7141 332e828 4ad7141 332e828 4ad7141 332e828 4ad7141 332e828 4ad7141 332e828 4ad7141 332e828 4ad7141 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
import gradio as gr
from transformers import Qwen3VLForConditionalGeneration, AutoProcessor
import torch
from PIL import Image
import spaces
import json
import re
# Load model and processor
model = Qwen3VLForConditionalGeneration.from_pretrained(
"Qwen/Qwen3-VL-2B-Instruct",
torch_dtype=torch.bfloat16,
device_map="auto"
)
processor = AutoProcessor.from_pretrained("Qwen/Qwen3-VL-2B-Instruct")
SYSTEM_PROMPT = """Eres un asistente que recibe la imagen de un ticket de compra y responde SOLO con JSON válido.
Esquema requerido:
{
"merchant": string,
"date": string | null,
"time": string | null,
"currency": string | null,
"subtotal": number | null,
"tax": number | null,
"total": number | null,
"paymentMethod": string | null,
"category": string | null,
"items": [
{ "name": string, "quantity": number | null, "unitPrice": number | null, "total": number | null }
]
}
Reglas:
- No inventes valores: si falta un dato, usa null.
- Los números deben ser numéricos, no strings.
- La salida debe ser SOLO ese JSON, sin texto extra ni bloques de código."""
@spaces.GPU(duration=120)
def analyze_ticket(image):
"""Analiza una imagen de ticket y devuelve JSON estructurado."""
if image is None:
return {"error": "No se proporcionó imagen"}
# Build message with image and prompt
messages = [
{
"role": "user",
"content": [
{"type": "image", "image": image},
{"type": "text", "text": SYSTEM_PROMPT + "\n\nAnaliza este ticket."}
]
}
]
# Prepare inputs using apply_chat_template
inputs = processor.apply_chat_template(
messages,
tokenize=True,
add_generation_prompt=True,
return_dict=True,
return_tensors="pt"
)
inputs = inputs.to(model.device)
# Generate response
with torch.no_grad():
generated_ids = model.generate(
**inputs,
max_new_tokens=1024,
do_sample=False,
)
# Decode output
generated_ids_trimmed = [
out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
]
output_text = processor.batch_decode(
generated_ids_trimmed,
skip_special_tokens=True,
clean_up_tokenization_spaces=False
)[0]
print(f"Respuesta del modelo: {output_text[:500]}")
# Parse JSON response
try:
cleaned = output_text.strip()
cleaned = re.sub(r'^```(?:json)?\s*', '', cleaned)
cleaned = re.sub(r'\s*```$', '', cleaned)
result = json.loads(cleaned)
return result
except json.JSONDecodeError as e:
print(f"Error parseando JSON: {e}")
return {"raw_response": output_text, "parse_error": True}
# Simple Interface with /predict endpoint
demo = gr.Interface(
fn=analyze_ticket,
inputs=gr.Image(type="pil", label="Imagen del ticket"),
outputs=gr.JSON(label="Datos extraídos"),
title="🧾 Ticket OCR",
description="Sube una imagen de un ticket para extraer la información estructurada.",
api_name="predict"
)
if __name__ == "__main__":
demo.launch()
|