Spaces:
Sleeping
Sleeping
File size: 4,088 Bytes
3658694 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
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()
|