| 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" |
|
|
| |
| 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." |
| |
| 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) |
| |
| if full_output.startswith(system_message): |
| summary = full_output[len(system_message):].strip() |
| else: |
| summary = full_output |
| |
| 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) |
