app.py

import gradio as gr
from arxiv2text import arxiv_to_text

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

def get_model(model_url="thepowerfuldeez/Qwen2-1.5B-Summarize"):
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B-Instruct")
    model = AutoModelForCausalLM.from_pretrained(
        model_url,
        bnb_4bit_compute_dtype=torch.bfloat16,
        load_in_4bit=True,
        attn_implementation="flash_attention_2",
    )
    return model, tokenizer
    

def call_llm(model, tokenizer, text):
    messages = [
        {"role": "system", "content": "You are helpful AI assistant."},
        {"role": "user", "content": text},
    ]
    input_ids = tokenizer.apply_chat_template(
        messages, add_generation_prompt=True, return_tensors="pt"
    )
    new_tokens = model.generate(input_ids, max_new_tokens=512)[0][len(input_ids[0]) :]
    output = tokenizer.decode(new_tokens, skip_special_tokens=True)
    return output

model, tokenizer = get_model()

def summarize_pdf(pdf_url):
    extracted_text = arxiv_to_text(pdf_url)
    summary = call_llm(model, tokenizer, f"Summarize following text: {extracted_text[:71000]}")
    return summary

interface = gr.Interface(
    fn=summarize_pdf,
    inputs="text",
    outputs="text",
    title="Arxiv PDF Summarizer",
    description="Enter the URL of an Arxiv PDF to get a summary."
)

interface.launch()