PDF-Summarizer / app.py
cstr's picture
Update app.py
2ec1f78 verified
raw
history blame
17.9 kB
import os
import re
import tempfile
import requests
import gradio as gr
from PyPDF2 import PdfReader
import logging
import webbrowser
from huggingface_hub import InferenceClient
from typing import Dict, List, Optional, Tuple
import time
# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# Constants
CONTEXT_SIZES = {
"4K": 4000,
"8K": 8000,
"32K": 32000,
"128K": 128000,
"200K": 200000
}
class ModelRegistry:
def __init__(self):
self.hf_models = {
"Phi-3 Mini 128k": "microsoft/Phi-3-mini-128k-instruct",
"Custom Model": ""
}
self.groq_models = self._fetch_groq_models()
def _fetch_groq_models(self) -> Dict[str, str]:
"""Fetch available Groq models with proper error handling"""
try:
groq_api_key = os.getenv('GROQ_API_KEY')
if not groq_api_key:
logging.warning("No GROQ_API_KEY found in environment")
return self._get_default_groq_models()
headers = {
"Authorization": f"Bearer {groq_api_key}",
"Content-Type": "application/json"
}
response = requests.get("https://api.groq.com/openai/v1/models", headers=headers)
if response.status_code == 200:
models = response.json().get("data", [])
return {model["id"]: model["id"] for model in models}
else:
logging.error(f"Failed to fetch Groq models: {response.status_code}")
return self._get_default_groq_models()
except Exception as e:
logging.error(f"Error fetching Groq models: {e}")
return self._get_default_groq_models()
def _get_default_groq_models(self) -> Dict[str, str]:
"""Return default Groq models when API is unavailable"""
return {
"llama-3.1-70b-versatile": "llama-3.1-70b-versatile",
"mixtral-8x7b-32768": "mixtral-8x7b-32768",
"llama-3.1-8b-instant": "llama-3.1-8b-instant"
}
def refresh_groq_models(self) -> Dict[str, str]:
"""Refresh the list of available Groq models"""
self.groq_models = self._fetch_groq_models()
return self.groq_models
# Initialize model registry
model_registry = ModelRegistry()
def extract_text_from_pdf(pdf_path: str) -> str:
"""Extract text content from PDF file."""
try:
reader = PdfReader(pdf_path)
text = ""
for page_num, page in enumerate(reader.pages, start=1):
page_text = page.extract_text()
if page_text:
text += page_text + "\n"
else:
logging.warning(f"No text found on page {page_num}.")
if not text.strip():
return "Error: No extractable text found in the PDF."
return text
except Exception as e:
logging.error(f"Error reading PDF file: {e}")
return f"Error reading PDF file: {e}"
def format_content(text: str, format_type: str) -> str:
"""Format extracted text according to specified format."""
if format_type == 'txt':
return text
elif format_type == 'md':
paragraphs = text.split('\n\n')
return '\n\n'.join(paragraphs)
elif format_type == 'html':
paragraphs = text.split('\n\n')
return ''.join([f'<p>{para.strip()}</p>' for para in paragraphs if para.strip()])
else:
logging.error(f"Unsupported format: {format_type}")
return f"Unsupported format: {format_type}"
def split_into_snippets(text: str, context_size: int) -> List[str]:
"""Split text into manageable snippets based on context size."""
sentences = re.split(r'(?<=[.!?]) +', text)
snippets = []
current_snippet = ""
for sentence in sentences:
if len(current_snippet) + len(sentence) + 1 > context_size:
if current_snippet:
snippets.append(current_snippet.strip())
current_snippet = sentence + " "
else:
snippets.append(sentence.strip())
current_snippet = ""
else:
current_snippet += sentence + " "
if current_snippet.strip():
snippets.append(current_snippet.strip())
return snippets
def build_prompts(snippets: List[str], prompt_instruction: str, custom_prompt: Optional[str], snippet_num: Optional[int] = None) -> str:
"""Build formatted prompts from text snippets."""
if snippet_num is not None:
if 1 <= snippet_num <= len(snippets):
selected_snippets = [snippets[snippet_num - 1]]
else:
return f"Error: Invalid snippet number. Please choose between 1 and {len(snippets)}."
else:
selected_snippets = snippets
prompts = []
base_prompt = custom_prompt if custom_prompt else prompt_instruction
for idx, snippet in enumerate(selected_snippets, start=1):
if len(selected_snippets) > 1:
prompt_header = f"{base_prompt} Part {idx} of {len(selected_snippets)}: ---\n"
else:
prompt_header = f"{base_prompt} ---\n"
framed_prompt = f"{prompt_header}{snippet}\n---"
prompts.append(framed_prompt)
return "\n\n".join(prompts)
def send_to_model(prompt, model_selection, hf_model_choice, hf_custom_model, hf_api_key,
groq_model_choice, groq_api_key, openai_api_key):
"""Send prompt to selected model"""
try:
if model_selection == "HuggingFace Inference":
if not hf_api_key:
return "HuggingFace API key required.", []
model_id = hf_custom_model if hf_model_choice == "Custom Model" else model_registry.hf_models[hf_model_choice]
summary = send_to_hf_inference(prompt, model_id, hf_api_key)
elif model_selection == "Groq API":
if not groq_api_key:
return "Groq API key required.", []
summary = send_to_groq(prompt, groq_model_choice, groq_api_key)
elif model_selection == "OpenAI ChatGPT":
if not openai_api_key:
return "OpenAI API key required.", []
# Implement OpenAI API call here
# Save summary for download
with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.txt') as summary_file:
summary_file.write(summary)
return summary, [summary_file.name]
except Exception as e:
logging.error(f"Error sending to model: {e}")
return f"Error sending to model: {str(e)}", []
def send_to_hf_inference(prompt: str, model_name: str, api_key: str) -> str:
"""Send prompt to HuggingFace using Inference API"""
try:
client = InferenceClient(api_key=api_key)
messages = [{"role": "user", "content": prompt}]
completion = client.chat.completions.create(
model=model_name,
messages=messages,
max_tokens=500
)
return completion.choices[0].message.content
except Exception as e:
logging.error(f"Error with HF inference: {e}")
return f"Error with HF inference: {e}"
def send_to_groq(prompt: str, model_name: str, api_key: str) -> str:
"""Send prompt to Groq API"""
try:
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
}
data = {
"model": model_name,
"messages": [{"role": "user", "content": prompt}]
}
response = requests.post(
"https://api.groq.com/openai/v1/chat/completions",
headers=headers,
json=data
)
return response.json()["choices"][0]["message"]["content"]
except Exception as e:
logging.error(f"Error with Groq API: {e}")
return f"Error with Groq API: {e}"
def copy_to_clipboard(text: str) -> str:
"""Copy text to clipboard using JavaScript"""
return """
navigator.clipboard.writeText(text)
.then(() => gradioApp().querySelector('#progress_status').value = 'Copied to clipboard!')
.catch(() => gradioApp().querySelector('#progress_status').value = 'Failed to copy');
"""
def open_chatgpt() -> None:
"""Open ChatGPT in new browser tab"""
return """window.open('https://chat.openai.com/', '_blank');"""
def process_pdf(pdf, fmt, ctx_size, snippet_num, prompt):
"""Generate prompt from PDF without model processing"""
try:
if not pdf:
return "Please upload a PDF file.", "", []
# Extract text
text = extract_text_from_pdf(pdf.name)
if text.startswith("Error"):
return text, "", []
# Format content
formatted_text = format_content(text, fmt)
# Split into snippets
snippets = split_into_snippets(formatted_text, ctx_size)
# Build prompts
default_prompt = "Summarize the following text:"
full_prompt = build_prompts(snippets, default_prompt, prompt, snippet_num)
if isinstance(full_prompt, str) and full_prompt.startswith("Error"):
return full_prompt, "", []
# Save prompt for download
with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.txt') as prompt_file:
prompt_file.write(full_prompt)
return "Prompt generated!", full_prompt, [prompt_file.name]
except Exception as e:
logging.error(f"Error processing PDF: {e}")
return f"Error processing PDF: {str(e)}", "", []
The main error is the context_size not being defined before it's used in the button click handlers. Let's fix the order of component definitions and handlers. Here's the corrected UI section:
python
Copy
# Main Interface
with gr.Blocks(theme=gr.themes.Default()) as demo:
# Store context size value
context_size_value = gr.State(value=32000)
# Header
gr.Markdown("# πŸ“„ Smart PDF Summarizer")
gr.Markdown("Upload a PDF document and get AI-powered summaries using various AI models.")
# Main Content
with gr.Row():
# Left Column - Input Options
with gr.Column(scale=1):
pdf_input = gr.File(
label="πŸ“ Upload PDF",
file_types=[".pdf"]
)
with gr.Row():
format_type = gr.Radio(
choices=["txt", "md", "html"],
value="txt",
label="πŸ“ Output Format"
)
# First define the slider
context_size = gr.Slider(
minimum=1000,
maximum=200000,
step=1000,
value=32000,
label="πŸ“ Custom Context Size"
)
# Then define the context size buttons
gr.Markdown("### Context Size")
with gr.Row():
for size_name, size_value in CONTEXT_SIZES.items():
gr.Button(
size_name,
size="sm", # Make buttons smaller
scale=1 # Equal scaling
).click(
lambda v=size_value: v,
None,
context_size
)
context_size = gr.Slider(
minimum=1000,
maximum=200000,
step=1000,
value=32000,
label="πŸ“ Custom Context Size"
)
snippet_number = gr.Number(
label="πŸ”’ Snippet Number",
value=1,
precision=0
)
custom_prompt = gr.Textbox(
label="✍️ Custom Prompt",
placeholder="Enter your custom prompt here...",
lines=2
)
model_choice = gr.Radio(
choices=["OpenAI ChatGPT", "HuggingFace Inference", "Groq API"],
value="OpenAI ChatGPT",
label="πŸ€– Model Selection"
)
with gr.Column(visible=False) as hf_options:
hf_model = gr.Dropdown(
choices=list(model_registry.hf_models.keys()),
label="πŸ”§ HuggingFace Model",
value="Phi-3 Mini 128k"
)
hf_custom_model = gr.Textbox(
label="Custom Model ID",
placeholder="Enter custom model ID...",
visible=False
)
hf_api_key = gr.Textbox(
label="πŸ”‘ HuggingFace API Key",
type="password"
)
with gr.Column(visible=False) as groq_options:
groq_model = gr.Dropdown(
choices=list(model_registry.groq_models.keys()),
label="πŸ”§ Groq Model",
value=list(model_registry.groq_models.keys())[0]
)
groq_refresh_btn = gr.Button("πŸ”„ Refresh Models")
groq_api_key = gr.Textbox(
label="πŸ”‘ Groq API Key",
type="password"
)
# In the UI section, add OpenAI API key input:
with gr.Column(visible=False) as openai_options:
openai_api_key = gr.Textbox(
label="πŸ”‘ OpenAI API Key",
type="password"
)
# Right Column - Output
with gr.Column(scale=1):
process_button = gr.Button("πŸš€ Process PDF", variant="primary")
progress_status = gr.Textbox(
label="πŸ“Š Progress",
interactive=False
)
generated_prompt = gr.Textbox(
label="πŸ“‹ Generated Prompt",
lines=10
)
with gr.Row():
copy_prompt_button = gr.Button("πŸ“‹ Copy Prompt")
open_chatgpt_button = gr.Button("🌐 Open ChatGPT")
summary_output = gr.Textbox(
label="πŸ“ Summary",
lines=15
)
with gr.Row():
copy_summary_button = gr.Button("πŸ“‹ Copy Summary")
download_files = gr.Files(
label="πŸ“₯ Download Files"
)
# Event Handlers
def update_context_size(size):
return gr.update(value=size)
def toggle_model_options(choice):
return (
gr.update(visible=choice == "HuggingFace Inference"),
gr.update(visible=choice == "Groq API"),
gr.update(visible=choice == "OpenAI ChatGPT")
)
def refresh_groq_models_list():
updated_models = model_registry.refresh_groq_models()
return gr.update(choices=list(updated_models.keys()))
def toggle_custom_model(model_name):
return gr.update(visible=model_name == "Custom Model")
# Connect event handlers
model_choice.change(
toggle_model_options,
inputs=[model_choice],
outputs=[hf_options, groq_options, openai_options]
)
for btn, size_value in context_buttons:
btn.click(
lambda v=size_value: v, # Simplified to directly return the value
None,
context_size
)
hf_model.change(
toggle_custom_model,
inputs=[hf_model],
outputs=[hf_custom_model]
)
groq_refresh_btn.click(
refresh_groq_models_list,
outputs=[groq_model]
)
process_button.click(
process_pdf,
inputs=[
pdf_input,
format_type,
context_size,
snippet_number,
custom_prompt
],
outputs=[
progress_status,
generated_prompt,
download_files
]
)
# Add a new button for sending to model
send_button = gr.Button("πŸš€ Send to Model", variant="primary")
send_button.click(
send_to_model,
inputs=[
generated_prompt,
model_choice,
hf_model,
hf_custom_model,
hf_api_key,
groq_model,
groq_api_key,
openai_api_key
],
outputs=[
summary_output,
download_files
]
)
copy_prompt_button.click(
copy_to_clipboard,
inputs=[generated_prompt],
outputs=[progress_status]
)
copy_summary_button.click(
copy_to_clipboard,
inputs=[summary_output],
outputs=[progress_status]
)
open_chatgpt_button.click(
open_chatgpt,
outputs=[progress_status]
)
# Instructions
gr.Markdown("""
### πŸ“Œ Instructions:
1. Upload a PDF document
2. Choose output format and context window size
3. Select snippet number (default: 1) or enter custom prompt
4. Select your preferred model in case you want to proceed directly (or continue with 5):
- OpenAI ChatGPT: Manual copy/paste workflow
- HuggingFace Inference: Direct API integration
- Groq API: High-performance inference
5. Click 'Process PDF' to generate summary
6. Use 'Copy Prompt' and, optionally, 'Open ChatGPT' for manual processing
7. Download generated files as needed
""")
# Launch the interface
if __name__ == "__main__":
demo.launch(share=False, debug=True)