File size: 3,427 Bytes
702ff78 0cf00a8 702ff78 4e7dc13 702ff78 5d660d7 702ff78 0dd885f 5d660d7 702ff78 5d660d7 7fce000 1eb3bd7 5d660d7 1eb3bd7 5d660d7 1eb3bd7 5d660d7 1eb3bd7 5d660d7 1eb3bd7 5d660d7 1eb3bd7 5d660d7 1eb3bd7 5d660d7 1eb3bd7 5d660d7 1eb3bd7 5d660d7 7fce000 5d660d7 | 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 | import os
os.system("pip install git+https://github.com/shumingma/transformers.git")
import threading
import torch
import torch._dynamo
torch._dynamo.config.suppress_errors = True
import numpy as np
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TextIteratorStreamer,
)
from sentence_transformers import SentenceTransformer
from PyPDF2 import PdfReader
from sklearn.metrics.pairwise import cosine_similarity
import gradio as gr
BITNET_ID = "microsoft/bitnet-b1.58-2B-4T"
# Load models
tokenizer = AutoTokenizer.from_pretrained(BITNET_ID)
bitnet = AutoModelForCausalLM.from_pretrained(
BITNET_ID,
torch_dtype=torch.bfloat16,
device_map="auto"
)
def extract_pdf_text(file_obj, max_chars=5000):
"""Extracts and returns up to max_chars of text from all pages of a PDF."""
reader = PdfReader(file_obj)
text = "\n".join([page.extract_text() for page in reader.pages if page.extract_text()])
return text[:max_chars]
def summarize_document(file, system_message, max_tokens, temperature, top_p):
if not file:
return "⚠️ Adjunta un archivo PDF primero."
# Extraer los primeros 5000 caracteres
text_fragment = extract_pdf_text(file, max_chars=5000)
if not text_fragment:
return "❌ No se pudo extraer texto del PDF."
input_text = system_message + "\n" + text_fragment
inputs = tokenizer(input_text, return_tensors="pt").to(bitnet.device)
output_ids = bitnet.generate(
**inputs,
max_new_tokens=max_tokens,
do_sample=True,
temperature=temperature,
top_p=top_p,
pad_token_id=tokenizer.eos_token_id
)
full_output = tokenizer.decode(output_ids[0], skip_special_tokens=True)
# Quitar el system_message duplicado, si lo hay
if full_output.startswith(system_message):
summary = full_output[len(system_message):].strip()
else:
summary = full_output
# Limitar la longitud del resumen a 1024 caracteres
max_output_length = 1024
if len(summary) > max_output_length:
summary = summary[:max_output_length] + "..."
return summary
css = """
.gradio-container { max-width: 800px !important; margin: auto; }
h1, h2 { text-align: center; }
"""
with gr.Blocks(css=css, theme=gr.themes.Soft(), title="Resumen BITNet de PDF Científico") as demo:
gr.HTML("<h1>🤖 Resumidor PDF con BITNet</h1><h2>Adjunta un PDF científico para recibir un resumen</h2>")
with gr.Row():
file = gr.File(label="Adjunta PDF", file_types=[".pdf"])
with gr.Column():
system_message = gr.Textbox(
value="You are a helpful assistant that summarizes scientific articles.",
label="System message"
)
max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens")
temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
summary = gr.Textbox(label="Resumen generado", lines=8)
summarize_btn = gr.Button("Generar resumen")
summarize_btn.click(
summarize_document,
inputs=[file, system_message, max_tokens, temperature, top_p],
outputs=summary
)
if __name__ == "__main__":
demo.launch(server_name="0.0.0.0", server_port=7860, share=True) |