Spaces:
Sleeping
Sleeping
File size: 1,651 Bytes
dd7f59d |
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 |
import torch
from transformers import T5TokenizerFast, T5ForConditionalGeneration
from torch.quantization import quantize_dynamic
import warnings
import gradio as gr
# Suprimir específicamente los avisos de TypedStorage
warnings.filterwarnings("ignore", category=UserWarning, message=".*TypedStorage is deprecated.*")
# Cargar el tokenizer y el modelo base
tokenizer = T5TokenizerFast.from_pretrained('google/flan-t5-small')
model = T5ForConditionalGeneration.from_pretrained('google/flan-t5-small')
model.eval() # Preparar el modelo para evaluación
# Aplicar cuantificación dinámica al modelo ANTES de cargar el estado del modelo cuantificado
quantized_model = quantize_dynamic(
model, {torch.nn.Linear}, dtype=torch.qint8
)
# Cargar el estado del modelo cuantificado
quantized_model.load_state_dict(torch.load('flan-t5-small-quantized.pth'))
# Función para resumir texto usando Flan-T5 cuantificado
def summarize(text):
inputs = tokenizer([text], max_length=1024, return_tensors="pt", truncation=True)
summary_ids = quantized_model.generate(inputs["input_ids"], num_beams=4, max_length=200, early_stopping=True)
summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
return summary
# Crear interfaz de Gradio
interface = gr.Interface(
fn=summarize,
inputs=gr.inputs.Textbox(lines=2, placeholder="Enter text to summarize..."),
outputs="text",
title="Text Summarizer",
description="A simple text summarizer based on the quantized Flan-T5 model."
)
# Lanzar la aplicación con Gradio, preparada para ser alojada en Hugging Face Spaces
if __name__ == "__main__":
interface.launch()
|