Spaces:
Sleeping
Sleeping
| import nbformat | |
| import gradio as gr | |
| from transformers import pipeline | |
| class NotebookEnhancer: | |
| def __init__(self): | |
| # Initialize Hugging Face models | |
| self.title_generator = pipeline( | |
| "summarization", model="facebook/bart-large-cnn" | |
| ) | |
| self.summary_generator = pipeline( | |
| "summarization", model="sshleifer/distilbart-cnn-12-6" | |
| ) | |
| def generate_title(self, code): | |
| """Generate a concise title for a code cell""" | |
| # Limit input length to match model constraints | |
| max_length = 1024 | |
| truncated_code = code[:max_length] if len(code) > max_length else code | |
| result = self.title_generator( | |
| truncated_code, max_length=10, min_length=3, do_sample=False | |
| ) | |
| title = result[0]["summary_text"].strip() | |
| # Format as a markdown title | |
| return f"## {title.capitalize()}" | |
| def generate_summary(self, code): | |
| """Generate a detailed summary for a code cell""" | |
| # Limit input length to match model constraints | |
| max_length = 1024 | |
| truncated_code = code[:max_length] if len(code) > max_length else code | |
| result = self.summary_generator( | |
| truncated_code, max_length=100, min_length=30, do_sample=True | |
| ) | |
| return result[0]["summary_text"].strip() | |
| def enhance_notebook(self, notebook_content): | |
| """Add title and summary markdown cells before each code cell""" | |
| # Load the notebook | |
| notebook = nbformat.reads(notebook_content, as_version=4) | |
| # Create a new notebook | |
| enhanced_notebook = nbformat.v4.new_notebook() | |
| enhanced_notebook.metadata = notebook.metadata | |
| # Process each cell | |
| i = 0 | |
| while i < len(notebook.cells): | |
| cell = notebook.cells[i] | |
| # For code cells, add title and summary markdown cells | |
| if cell.cell_type == "code" and cell.source.strip(): | |
| # Generate title | |
| title = self.generate_title(cell.source) | |
| title_cell = nbformat.v4.new_markdown_cell(title) | |
| enhanced_notebook.cells.append(title_cell) | |
| # Generate summary | |
| summary = self.generate_summary(cell.source) | |
| summary_cell = nbformat.v4.new_markdown_cell(summary) | |
| enhanced_notebook.cells.append(summary_cell) | |
| # Add the original cell | |
| enhanced_notebook.cells.append(cell) | |
| i += 1 | |
| # Convert back to string | |
| return nbformat.writes(enhanced_notebook) | |
| def process_notebook(file): | |
| """Process an uploaded notebook file""" | |
| enhancer = NotebookEnhancer() | |
| # Read uploaded file | |
| notebook_content = file.decode("utf-8") | |
| # Process the notebook | |
| enhanced_notebook = enhancer.enhance_notebook(notebook_content) | |
| # Save to temp file | |
| output_path = "enhanced_notebook.ipynb" | |
| with open(output_path, "w", encoding="utf-8") as f: | |
| f.write(enhanced_notebook) | |
| return output_path | |
| def build_gradio_interface(): | |
| """Create and launch the Gradio interface""" | |
| with gr.Blocks(title="Notebook Enhancer") as demo: | |
| gr.Markdown("# Jupyter Notebook Enhancer") | |
| gr.Markdown( | |
| """ | |
| Upload a Jupyter notebook to enhance it with automatically generated titles and summaries for each code cell. | |
| This tool uses Hugging Face models to: | |
| 1. Generate concise titles for code cells | |
| 2. Create explanatory summaries of what the code does | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| file_input = gr.File(label="Upload Jupyter Notebook (.ipynb)") | |
| process_btn = gr.Button("Enhance Notebook") | |
| with gr.Column(): | |
| output = gr.File(label="Enhanced Notebook") | |
| process_btn.click(fn=process_notebook, inputs=file_input, outputs=output) | |
| return demo | |
| # This will be the entry point when running the script | |
| if __name__ == "__main__": | |
| demo = build_gradio_interface() | |
| demo.launch() | |