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)