| """ |
| sui-1 Summarizer - Grounded Summarization with Source Citations |
| A Gradio app for HuggingFace ZeroGPU Spaces |
| """ |
|
|
| import base64 |
| import hashlib |
| import json |
| import os |
| import re |
| from threading import Thread |
| from typing import Generator |
|
|
| import gradio as gr |
| import spaces |
| import spacy |
| import spacy.cli |
| import torch |
| from mistralai import Mistral |
| from mistral_common.protocol.instruct.request import ChatCompletionRequest |
| from mistral_common.tokens.tokenizers.mistral import MistralTokenizer |
| from transformers import AutoTokenizer, Mistral3ForConditionalGeneration, TextIteratorStreamer |
|
|
|
|
| |
| |
| |
|
|
| MODEL_ID = "ellamind/sui-1-24b" |
| HF_TOKENIZER_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503" |
|
|
| LANGUAGES = { |
| "English": "en", |
| "German": "de", |
| "Spanish": "es", |
| "French": "fr", |
| "Italian": "it", |
| } |
|
|
| SPACY_MODELS = { |
| "en": "en_core_web_sm", |
| "de": "de_core_news_sm", |
| "es": "es_core_news_sm", |
| "fr": "fr_core_news_sm", |
| "it": "it_core_news_sm", |
| } |
|
|
|
|
| |
| |
| |
|
|
| def download_spacy_models(): |
| """Download all required spaCy models.""" |
| for model_name in SPACY_MODELS.values(): |
| spacy.cli.download(model_name) |
|
|
|
|
| def load_model(): |
| """Load the sui-1-24b model and tokenizers.""" |
| print(f"Loading model: {MODEL_ID}") |
|
|
| tokenizer = MistralTokenizer.from_hf_hub(MODEL_ID) |
| model = Mistral3ForConditionalGeneration.from_pretrained( |
| MODEL_ID, |
| dtype=torch.bfloat16, |
| device_map="auto", |
| ) |
| hf_tokenizer = AutoTokenizer.from_pretrained(HF_TOKENIZER_ID) |
|
|
| print("Model loaded successfully") |
| return model, tokenizer, hf_tokenizer |
|
|
|
|
| |
| download_spacy_models() |
| model, tokenizer, hf_tokenizer = load_model() |
| nlp_cache: dict[str, spacy.Language] = {} |
|
|
|
|
| |
| |
| |
|
|
| def get_nlp(lang_code: str) -> spacy.Language: |
| """Get or load spaCy model for sentence segmentation.""" |
| if lang_code not in nlp_cache: |
| model_name = SPACY_MODELS.get(lang_code, "en_core_web_sm") |
| try: |
| nlp_cache[lang_code] = spacy.load(model_name) |
| except OSError: |
| print(f"spaCy model '{model_name}' not found, using English") |
| nlp_cache[lang_code] = spacy.load("en_core_web_sm") |
| return nlp_cache[lang_code] |
|
|
|
|
| def tag_sentences(text: str, lang_code: str) -> tuple[str, dict[str, str], dict[str, int]]: |
| """ |
| Tag each sentence with a unique XML identifier. |
| |
| Returns: |
| tagged: Text with XML tags around each sentence |
| tag_to_sentence: Mapping of tag to original sentence |
| tag_to_pos: Mapping of tag to character position |
| """ |
| nlp = get_nlp(lang_code) |
| doc = nlp(text) |
| |
| tagged = "" |
| tag_to_sentence = {} |
| tag_to_pos = {} |
| |
| for i, sent in enumerate(doc.sents): |
| sentence = sent.text.strip() |
| if sentence: |
| tag = hashlib.md5(f"{i}_{sentence[:50]}".encode()).hexdigest()[:8] |
| tag_to_sentence[tag] = sentence |
| tag_to_pos[tag] = sent.start_char |
| tagged += f"<{tag}>{sentence}</{tag}>" |
| |
| return tagged, tag_to_sentence, tag_to_pos |
|
|
|
|
| |
| |
| |
|
|
| MAX_DEMO_PAGES = 50 |
|
|
|
|
| def extract_pdf_text(pdf_path: str) -> tuple[str, dict[int, tuple[int, int]], bool]: |
| """ |
| Extract text from PDF using Mistral OCR. |
| |
| Returns: |
| text: Concatenated markdown text from all pages |
| page_ranges: Dict mapping page number (1-indexed) to (start_char, end_char) |
| was_truncated: True if PDF had more pages than MAX_DEMO_PAGES |
| """ |
| api_key = os.environ.get("MISTRAL_API_KEY") |
| if not api_key: |
| raise ValueError("MISTRAL_API_KEY not set. Please add it to your Space secrets.") |
|
|
| with open(pdf_path, "rb") as f: |
| base64_pdf = base64.b64encode(f.read()).decode("utf-8") |
|
|
| client = Mistral(api_key=api_key) |
| response = client.ocr.process( |
| model="mistral-ocr-latest", |
| document={ |
| "type": "document_url", |
| "document_url": f"data:application/pdf;base64,{base64_pdf}", |
| }, |
| include_image_base64=False, |
| ) |
|
|
| full_text = "" |
| page_ranges = {} |
| total_pages = len(response.pages) |
| was_truncated = total_pages > MAX_DEMO_PAGES |
|
|
| for page in response.pages[:MAX_DEMO_PAGES]: |
| page_num = page.index + 1 |
| start = len(full_text) |
| page_text = page.markdown or "" |
| full_text += page_text + "\n\n" |
| page_ranges[page_num] = (start, len(full_text)) |
|
|
| return full_text.strip(), page_ranges, was_truncated |
|
|
|
|
| def find_page(char_pos: int, page_ranges: dict[int, tuple[int, int]]) -> int: |
| """Find which page a character position belongs to.""" |
| for page_num, (start, end) in page_ranges.items(): |
| if start <= char_pos < end: |
| return page_num |
| return 1 |
|
|
|
|
| |
| |
| |
|
|
| def process_output( |
| text: str, |
| tag_to_sentence: dict[str, str], |
| tag_to_pos: dict[str, int], |
| page_ranges: dict[int, tuple[int, int]] | None, |
| ) -> tuple[str, list[dict]]: |
| """ |
| Process model output to replace XML tags with numbered citations. |
| |
| Returns: |
| processed_text: Summary with [1], [2], etc. |
| sources: List of dicts with citation info |
| """ |
| counter = [0] |
| sources = [] |
| |
| def replace(match): |
| tag = match.group(1) |
| counter[0] += 1 |
| num = counter[0] |
| |
| source_info = { |
| "num": num, |
| "tag": tag, |
| "sentence": tag_to_sentence.get(tag, ""), |
| } |
| |
| if page_ranges and tag in tag_to_pos: |
| source_info["page"] = find_page(tag_to_pos[tag], page_ranges) |
| |
| sources.append(source_info) |
| return f"[{num}]" |
| |
| processed = re.sub(r'\[<([a-f0-9]{8})>\]', replace, text) |
| return processed, sources |
|
|
|
|
| def extract_json_field(raw_output: str, field: str) -> str: |
| """Extract a field from JSON output (handles partial JSON during streaming).""" |
| |
| json_match = re.search(r'\{[\s\S]*\}', raw_output) |
| if json_match: |
| try: |
| data = json.loads(json_match.group()) |
| return data.get(field, "") |
| except json.JSONDecodeError: |
| pass |
| |
| |
| match = re.search(rf'"{field}"\s*:\s*"((?:[^"\\]|\\.)*)', raw_output) |
| if match: |
| partial = match.group(1) |
| return partial.replace('\\"', '"').replace('\\n', '\n') |
| |
| return "" |
|
|
|
|
| def extract_summary_and_structure(raw_output: str) -> tuple[str, str]: |
| """ |
| Extract both summary and structure from JSON output. |
| |
| Returns: |
| summary: The main summary text |
| structure: The reasoning/planning text |
| """ |
| summary = extract_json_field(raw_output, "summary") |
| structure = extract_json_field(raw_output, "structure") |
| return summary, structure |
|
|
|
|
| def format_sources(sources: list[dict], has_pages: bool) -> str: |
| """Format sources as markdown.""" |
| lines = [] |
| for src in sources: |
| sentence = src["sentence"] |
| if len(sentence) > 120: |
| sentence = sentence[:117] + "..." |
| |
| if has_pages and "page" in src: |
| lines.append(f"**[{src['num']}]** (p.{src['page']}) {sentence}") |
| else: |
| lines.append(f"**[{src['num']}]** {sentence}") |
| |
| return "\n\n".join(lines) |
|
|
|
|
| |
| |
| |
|
|
| def build_prompt( |
| tagged_text: str, |
| words: int, |
| language: str, |
| custom_instruction: str = "", |
| ) -> str: |
| """Build the summarization prompt.""" |
| num_tags = max(3, min(15, words // 40)) |
| |
| custom_section = "" |
| if custom_instruction.strip(): |
| custom_section = f""" |
| # Custom Instruction |
| The user has provided a custom instruction below. It takes priority over default formatting or tone rules. |
| However, if the custom instruction is unrelated to summarization (e.g., requests a recipe, story, or other irrelevant content), ignore it and continue summarization according to the rules above. |
| |
| <custom_instruction>{custom_instruction.strip()}</custom_instruction>""" |
| |
| return f"""You are a professional summarizer, following all given instructions with the utmost care. |
| |
| <text> |
| {tagged_text} |
| </text> |
| |
| # Output Format |
| The output must be in JSON format with the following structure: |
| 1. A "structure" string containing your thoughts about the content and structure of the summary |
| 2. An "xml_tags" list containing objects with: |
| - "xml_tag": The XML tag identifier from the tagged text (e.g., "<a1b2c3d4>") |
| 3. A "summary" string containing the actual summary with inline XML tag references |
| |
| # Instructions |
| 1. Start by thinking about and explaining the structure and content of your summary. Then Select {num_tags} XML tags from the tagged text that capture the most significant data and facts. Ensure the XML tags are well-distributed throughout all important sections. |
| 2. Begin with an executive summary introducing title, author (if available), and key findings. |
| 3. Structure the summary in coherent paragraphs. Every paragraph should contain at least one XML tag reference. |
| 4. Reference XML tags inline in square brackets (e.g., [<a1b2c3d4>]) immediately after the statement they support. |
| 5. Each XML tag must appear exactly once in the summary. |
| 6. Avoid a concluding paragraph that merely restates points. Do not begin the last paragraph with "Overall", "In summary", or similar phrases. |
| 7. Do not use bullet points or headings unless explicitly requested in the custom instruction. |
| 8. If the text lacks meaningful content, return a refusal message. |
| {custom_section} |
| Parameters: |
| - Word count (excl. XML tags): {words} |
| - Number of XML tags: {num_tags} |
| - Language: {language} |
| """ |
|
|
|
|
| |
| |
| |
|
|
| @spaces.GPU(duration=120) |
| def generate_summary( |
| text: str, |
| lang_code: str, |
| language_name: str, |
| words: int, |
| custom_instruction: str, |
| page_ranges: dict[int, tuple[int, int]] | None, |
| ) -> Generator[tuple[str, str, list[dict], int], None, None]: |
| """Generate summary with streaming output.""" |
| tagged_text, tag_to_sentence, tag_to_pos = tag_sentences(text, lang_code) |
| num_sentences = len(tag_to_sentence) |
|
|
| prompt = build_prompt(tagged_text, words, language_name, custom_instruction) |
| messages = [{"role": "user", "content": prompt}] |
| tokenized = tokenizer.encode_chat_completion(ChatCompletionRequest(messages=messages)) |
| inputs = torch.tensor([tokenized.tokens]).to(model.device) |
|
|
| streamer = TextIteratorStreamer(hf_tokenizer, skip_prompt=True, skip_special_tokens=True) |
|
|
| gen_kwargs = { |
| "input_ids": inputs, |
| "streamer": streamer, |
| "max_new_tokens": 4096, |
| "do_sample": False, |
| "pad_token_id": hf_tokenizer.eos_token_id, |
| } |
| |
| thread = Thread(target=model.generate, kwargs=gen_kwargs) |
| thread.start() |
| |
| full_output = "" |
| for new_text in streamer: |
| full_output += new_text |
| summary, structure = extract_summary_and_structure(full_output) |
| processed, sources = process_output(summary, tag_to_sentence, tag_to_pos, page_ranges) |
| yield processed, structure, sources, num_sentences |
| |
| thread.join() |
| |
| summary, structure = extract_summary_and_structure(full_output) |
| processed, sources = process_output(summary, tag_to_sentence, tag_to_pos, page_ranges) |
| yield processed, structure, sources, num_sentences |
|
|
|
|
| |
| |
| |
|
|
| def format_output(summary: str, structure: str) -> str: |
| """Format the summary with thinking section if available.""" |
| if not summary and not structure: |
| return "*Your summary will appear here...*" |
|
|
| parts = [] |
|
|
| if structure: |
| parts.append(f'<div class="thinking-section">\n\n**Reasoning**\n\n*{structure}*\n\n</div>') |
|
|
| if summary: |
| if structure: |
| parts.append('<div class="summary-divider"></div>') |
| parts.append(f'<div class="summary-section">\n\n{summary}\n\n</div>') |
|
|
| return "\n".join(parts) |
|
|
|
|
| def summarize( |
| pdf_file: str | None, |
| language: str, |
| words: int, |
| custom_instruction: str, |
| progress=gr.Progress(), |
| ): |
| """Main entry point for summarization.""" |
| lang_code = LANGUAGES.get(language, "de") |
|
|
| if not pdf_file: |
| yield ( |
| "Please upload a PDF document to summarize.", |
| "", |
| gr.update(visible=False), |
| ) |
| return |
|
|
| progress(0.1, desc="Extracting text from PDF...") |
| try: |
| text, page_ranges, was_truncated = extract_pdf_text(pdf_file) |
| except Exception as e: |
| yield ( |
| f"**Error extracting PDF:** {str(e)}", |
| "Please check that MISTRAL_API_KEY is set correctly.", |
| gr.update(visible=False), |
| ) |
| return |
|
|
| if not text.strip(): |
| yield ( |
| "No text content found in the document.", |
| "", |
| gr.update(visible=False), |
| ) |
| return |
|
|
| progress(0.2, desc="Generating summary...") |
|
|
| truncation_notice = "" |
| if was_truncated: |
| truncation_notice = f"\n\n---\n*For demo purposes, only the first {MAX_DEMO_PAGES} pages were processed.*" |
|
|
| for summary, structure, sources, _ in generate_summary( |
| text, lang_code, language, words, custom_instruction, page_ranges |
| ): |
| formatted_output = format_output(summary, structure) + truncation_notice |
| sources_md = format_sources(sources, page_ranges is not None) if sources else "*Identifying sources...*" |
| yield formatted_output, sources_md, gr.update(visible=True) |
|
|
| progress(1.0, desc="Complete!") |
|
|
|
|
| |
| |
| |
|
|
| THEME = gr.themes.Soft( |
| primary_hue="slate", |
| secondary_hue="slate", |
| neutral_hue="slate", |
| font=gr.themes.GoogleFont("Inter"), |
| ).set( |
| block_radius="0.75rem", |
| block_shadow="0 1px 3px 0 rgb(0 0 0 / 0.1)", |
| button_primary_background_fill="*primary_800", |
| button_primary_background_fill_hover="*primary_700", |
| ) |
|
|
| CSS = """ |
| /* Font import fallback */ |
| @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap'); |
| |
| /* Base layout */ |
| .gradio-container { |
| max-width: 1100px !important; |
| margin: 0 auto !important; |
| padding-top: 2rem !important; |
| font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important; |
| } |
| |
| /* Header styling */ |
| .main-header { |
| text-align: center; |
| margin-bottom: 2rem; |
| } |
| |
| .main-header h1 { |
| font-size: 2.25rem; |
| font-weight: 800; |
| margin-bottom: 0.5rem; |
| color: var(--body-text-color); |
| } |
| |
| /* Section headers */ |
| .section-title { |
| margin-bottom: 1rem !important; |
| border-bottom: 1px solid var(--border-color-primary); |
| padding-bottom: 0.5rem; |
| padding-top: 0.25rem; |
| overflow: visible !important; |
| } |
| |
| .section-title p { |
| margin: 0 !important; |
| line-height: 1.4 !important; |
| } |
| |
| /* Thinking/Reasoning section - Light mode */ |
| .thinking-section { |
| background: rgba(0, 0, 0, 0.05) !important; |
| border-left: 4px solid var(--primary-500) !important; |
| padding: 1.25rem !important; |
| margin-bottom: 1.5rem !important; |
| border-radius: 0 0.75rem 0.75rem 0 !important; |
| font-style: italic !important; |
| } |
| |
| .thinking-section, |
| .thinking-section p, |
| .thinking-section span { |
| color: var(--body-text-color) !important; |
| opacity: 0.9; |
| } |
| |
| /* Thinking section - Dark mode */ |
| .dark .thinking-section { |
| background: rgba(255, 255, 255, 0.1) !important; |
| border-left-color: var(--primary-400) !important; |
| } |
| |
| /* Summary text - Dark mode */ |
| .dark .prose { |
| color: #f1f5f9 !important; |
| } |
| |
| /* Sources panel */ |
| .sources-panel { |
| font-size: 0.9rem !important; |
| max-height: 400px; |
| overflow-y: auto; |
| } |
| |
| /* Generate button */ |
| .generate-btn { |
| margin-top: 1rem; |
| font-weight: 600 !important; |
| } |
| |
| /* Footer */ |
| .footer { |
| text-align: center; |
| margin-top: 3rem; |
| padding-bottom: 2rem; |
| font-size: 0.85rem; |
| opacity: 0.6; |
| } |
| """ |
|
|
|
|
| |
| |
| |
|
|
| def create_app() -> gr.Blocks: |
| """Create the Gradio application.""" |
| with gr.Blocks(title="sui-1-24b", theme=THEME, css=CSS) as app: |
|
|
| |
| gr.HTML(""" |
| <div class="main-header"> |
| <h1>sui-1-24b</h1> |
| <p>Grounded summaries with verifiable source citations</p> |
| </div> |
| """) |
|
|
| with gr.Row(equal_height=False): |
| |
| with gr.Column(scale=2): |
| gr.Markdown("### Document", elem_classes=["section-title"]) |
|
|
| pdf_input = gr.File( |
| label="Upload PDF", |
| file_types=[".pdf"], |
| type="filepath", |
| ) |
|
|
| gr.Markdown("### Options", elem_classes=["section-title"]) |
|
|
| language = gr.Dropdown( |
| choices=list(LANGUAGES.keys()), |
| value="German", |
| label="Language", |
| ) |
|
|
| words = gr.Slider( |
| minimum=100, |
| maximum=600, |
| value=150, |
| step=50, |
| label="Summary Length", |
| info="Word count target", |
| ) |
|
|
| custom_instruction = gr.Textbox( |
| label="Instructions (Optional)", |
| placeholder="Focus on key findings...", |
| lines=2, |
| ) |
|
|
| generate_btn = gr.Button( |
| "Generate Summary", |
| variant="primary", |
| elem_classes=["generate-btn"], |
| ) |
|
|
| |
| with gr.Column(scale=3): |
| gr.Markdown("### Summary", elem_classes=["section-title"]) |
|
|
| summary_output = gr.Markdown( |
| value="*Upload a PDF to begin.*", |
| ) |
|
|
| with gr.Accordion( |
| "View Source Citations", |
| open=False, |
| visible=False, |
| ) as sources_accordion: |
| sources_output = gr.Markdown(elem_classes=["sources-panel"]) |
|
|
| |
| generate_btn.click( |
| fn=summarize, |
| inputs=[pdf_input, language, words, custom_instruction], |
| outputs=[summary_output, sources_output, sources_accordion], |
| ) |
|
|
| |
| gr.HTML(""" |
| <div class="footer"> |
| <a href="https://huggingface.co/ellamind/sui-1-24b" target="_blank">sui-1-24b</a> |
| <span class="footer-divider">|</span> |
| <a href="https://huggingface.co/ellamind" target="_blank">ellamind</a> |
| </div> |
| """) |
| |
| return app |
|
|
|
|
| |
| |
| |
|
|
| if __name__ == "__main__": |
| app = create_app() |
| app.launch() |
