Spaces:
Build error
Build error
| import os | |
| import sys | |
| import subprocess | |
| import re | |
| from collections.abc import Iterator | |
| import gradio as gr | |
| from huggingface_hub import hf_hub_download | |
| # Install llama-cpp-python if not present | |
| try: | |
| from llama_cpp import Llama | |
| except ModuleNotFoundError: | |
| subprocess.check_call([sys.executable, "-m", "pip", "install", "llama-cpp-python"]) | |
| from llama_cpp import Llama | |
| # Install requests if not present | |
| try: | |
| import requests | |
| except ModuleNotFoundError: | |
| subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"]) | |
| import requests | |
| # Install pypdf if not present | |
| try: | |
| from pypdf import PdfReader | |
| except ModuleNotFoundError: | |
| subprocess.check_call([sys.executable, "-m", "pip", "install", "pypdf"]) | |
| from pypdf import PdfReader | |
| MAX_MAX_NEW_TOKENS = 512 # Adjusted for document summaries | |
| DEFAULT_MAX_NEW_TOKENS = 256 | |
| MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096")) # Increased for longer context | |
| DESCRIPTION = """ | |
| # DocChat: PDF Summarizer and Q&A (CPU-Only, Ultra-Fast Optimization) | |
| This application delivers an interactive chat interface powered by a highly efficient, small AI model adapted for summarizing and answering questions about any publicly accessible PDF document through specialized prompt engineering. Provide a PDF URL in your message (e.g., 'Load this PDF: https://example.com/document.pdf'), and the bot will download, process it, and be ready for summaries or questions. | |
| It ensures rapid, reasoned responses to user queries. Duplicate this Space for customization or queue-free deployment. | |
| 🔎 Model details are available at the [google/gemma-2b-it](https://huggingface.co/google/gemma-2b-it) repository; pre-trained and instruction-tuned, further adapted here for document summarization. | |
| <p>Running on CPU 🥶 Inference is heavily optimized for responses in under 10 seconds for simple queries, with output limited to 512 tokens maximum. For longer, more complete responses, increase 'Max New Tokens' in Advanced Settings. Brief delays may occur in free-tier environments due to shared resources, but typical generation speeds reach 20-40 tokens per second.</p> | |
| """ | |
| LICENSE = """ | |
| <p/> | |
| --- | |
| This application employs the [google/gemma-2b-it](https://huggingface.co/google/gemma-2b-it) model, governed by Google's Gemma Terms of Use. Refer to the [model card](https://huggingface.co/google/gemma-2b-it) and [Gemma documentation](https://ai.google.dev/gemma/terms) for details. | |
| """ | |
| # Download the GGUF model file | |
| model_repo = "mradermacher/gemma-2b-it-GGUF" # Adjusted to match common repo; update if needed | |
| model_filename = "gemma-2b-it.Q4_K_M.gguf" | |
| model_path = hf_hub_download(repo_id=model_repo, filename=model_filename) | |
| # Load the model with optimizations | |
| llm = Llama( | |
| model_path=model_path, | |
| n_ctx=8192, # Increased context window for document content | |
| n_batch=512, # Balanced for CPU efficiency | |
| n_threads=2, # Set to 2 for Hugging Face free tier (2 vCPUs) | |
| n_gpu_layers=0 # Enforce CPU-only execution | |
| ) | |
| DEFAULT_SYSTEM_PROMPT = """You are DocChat, a knowledgeable AI assistant specializing in summarizing and analyzing PDF documents. Provide accurate, helpful, reasoned, detailed, and comprehensive summaries or answers based on the provided document. Always base responses on the document content. If no document is loaded, ask the user to provide a publicly accessible PDF URL.""" | |
| def process_pdf(url: str) -> str: | |
| try: | |
| response = requests.get(url, timeout=10) | |
| response.raise_for_status() | |
| pdf_path = "temp.pdf" | |
| with open(pdf_path, "wb") as f: | |
| f.write(response.content) | |
| reader = PdfReader(pdf_path) | |
| document_text = "" | |
| for page in reader.pages: | |
| text = page.extract_text() | |
| if text: | |
| document_text += text + "\n\n" | |
| # Truncate document text if too long | |
| max_doc_tokens = 6000 # Leave room for prompt and response | |
| doc_tokens = llm.tokenize(document_text.encode("utf-8")) | |
| if len(doc_tokens) > max_doc_tokens: | |
| document_text = llm.detokenize(doc_tokens[:max_doc_tokens]).decode("utf-8", errors="ignore") + "\n[Document truncated for context length]" | |
| os.remove(pdf_path) | |
| return document_text | |
| except Exception as e: | |
| return f"Error processing PDF: {str(e)}" | |
| def generate( | |
| message: str, | |
| chat_history: list[tuple[str, str]], | |
| system_prompt: str, | |
| max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS, | |
| temperature: float = 0.6, | |
| top_p: float = 0.9, | |
| top_k: int = 50, | |
| repetition_penalty: float = 1.2, | |
| ) -> Iterator[str]: | |
| if message.lower().strip() == "hi": | |
| response = "Hello, I am your document summarizer assistant. Provide a publicly accessible PDF URL to get started." | |
| yield response | |
| return | |
| # Build conversation for Gemma prompt format | |
| conversation = [] | |
| if system_prompt and not chat_history: | |
| # Prepend system prompt to first user message for Gemma | |
| first_message = system_prompt + "\n\n" + message | |
| conversation.append({"role": "user", "content": first_message}) | |
| else: | |
| if system_prompt: | |
| conversation.append({"role": "user", "content": system_prompt}) | |
| for user, assistant in chat_history: | |
| conversation.append({"role": "user", "content": user}) | |
| conversation.append({"role": "model", "content": assistant}) | |
| conversation.append({"role": "user", "content": message}) | |
| # Construct prompt in Gemma format | |
| prompt_parts = ["<bos>"] # Start with BOS token for Gemma | |
| for dialogue in conversation: | |
| role = "user" if dialogue["role"] == "user" else "model" | |
| prompt_parts.append(f"<start_of_turn>{role}\n{dialogue['content']}<end_of_turn>\n") | |
| prompt_parts.append("<start_of_turn>model\n") | |
| prompt = "".join(prompt_parts) | |
| # Tokenize and check length | |
| input_tokens = llm.tokenize(prompt.encode("utf-8"), add_bos=False) # Gemma handles BOS internally if needed | |
| if len(input_tokens) > MAX_INPUT_TOKEN_LENGTH: | |
| yield "Error: Input too long. Please shorten your query or simplify the request." | |
| return | |
| # Generate response | |
| response = "" | |
| tokens_generated = 0 | |
| stop_appending = False | |
| for token in llm.generate( | |
| input_tokens, | |
| top_k=top_k, | |
| top_p=top_p, | |
| temp=temperature, | |
| repeat_penalty=repetition_penalty, | |
| ): | |
| detokenized = llm.detokenize([token]).decode("utf-8", errors="ignore") | |
| if not stop_appending: | |
| response += detokenized | |
| if "<end_of_turn>" in detokenized or token == llm.token_eos(): | |
| stop_appending = True | |
| response = response.split("<end_of_turn>")[0].strip() # Clean up | |
| yield response | |
| tokens_generated += 1 | |
| if tokens_generated >= max_new_tokens: | |
| break | |
| # Gradio interface setup | |
| with gr.Blocks(css="""#chatbot {height: 500px;}""") as demo: | |
| gr.Markdown(DESCRIPTION) | |
| chatbot = gr.Chatbot(label="DocChat") | |
| msg = gr.Textbox(label="Provide PDF URL or ask about the PDF", placeholder="e.g., 'Load PDF: https://example.com/document.pdf' or 'Summarize the key points'") | |
| with gr.Row(): | |
| submit = gr.Button("Submit") | |
| clear = gr.Button("Clear") | |
| with gr.Accordion("Advanced Settings", open=False): | |
| system_prompt = gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, lines=6) | |
| temperature = gr.Slider(label="Temperature", value=0.6, minimum=0.0, maximum=1.0, step=0.05) | |
| top_p = gr.Slider(label="Top P", value=0.9, minimum=0.0, maximum=1.0, step=0.05) | |
| top_k = gr.Slider(label="Top K", value=50, minimum=1, maximum=100, step=1) | |
| repetition_penalty = gr.Slider(label="Repetition Penalty", value=1.2, minimum=1.0, maximum=2.0, step=0.05) | |
| max_new_tokens = gr.Slider(label="Max New Tokens", value=DEFAULT_MAX_NEW_TOKENS, minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1) | |
| gr.Markdown(LICENSE) | |
| def user(message, history): | |
| return "", history + [[message, None]] | |
| def bot(history, sys_prompt, temp, tp, tk, rp, mnt): | |
| message = history[-1][0] | |
| history[-1][1] = "" | |
| # Detect PDF URL in message | |
| pdf_url_match = re.search(r'(?i)https?://\S+\.pdf', message) | |
| new_sys_prompt = sys_prompt | |
| if pdf_url_match: | |
| url = pdf_url_match.group(0) | |
| document_text = process_pdf(url) | |
| if document_text.startswith("Error"): | |
| history[-1][1] = document_text | |
| yield history, new_sys_prompt | |
| return | |
| new_sys_prompt = f"""You are DocChat, a knowledgeable AI assistant specializing in summarizing and analyzing PDF documents. Provide accurate, helpful, reasoned, detailed, and comprehensive summaries or answers based on the provided document. Always base responses on the document content. | |
| The document is from {url}. Content excerpt: | |
| {document_text} | |
| Example: | |
| User: Summarize the key points of the document. | |
| Assistant: [Your summary based on the content].""" | |
| history[-1][1] = "PDF loaded and processed successfully. You can now ask questions or request a summary." | |
| yield history, new_sys_prompt | |
| return | |
| # If no PDF loaded and not loading, remind | |
| if "Content excerpt:" not in sys_prompt: | |
| history[-1][1] = "Please provide a publicly accessible PDF URL to load a document first." | |
| yield history, new_sys_prompt | |
| return | |
| # Generate response for queries | |
| for new_text in generate(message, history[:-1], sys_prompt, mnt, temp, tp, tk, rp): | |
| history[-1][1] = new_text | |
| yield history, new_sys_prompt | |
| submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then( | |
| bot, [chatbot, system_prompt, temperature, top_p, top_k, repetition_penalty, max_new_tokens], [chatbot, system_prompt] | |
| ) | |
| msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then( | |
| bot, [chatbot, system_prompt, temperature, top_p, top_k, repetition_penalty, max_new_tokens], [chatbot, system_prompt] | |
| ) | |
| clear.click(lambda: None, None, chatbot, queue=False) | |
| demo.queue(max_size=128).launch() |