app.py

import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import re
import time

# Model constants
CODET5_MODEL = "Salesforce/codet5-base-multi-sum"

class CodeT5Summarizer:
    def __init__(self, device=None):
        """Initialize CodeT5 summarization model."""
        self.device = device if device else ('cuda' if torch.cuda.is_available() else 'cpu')
        
        # Initialize model and tokenizer
        with st.spinner("Loading CodeT5 model... this may take a minute..."):
            self.tokenizer = AutoTokenizer.from_pretrained(CODET5_MODEL)
            self.model = AutoModelForSeq2SeqLM.from_pretrained(CODET5_MODEL).to(self.device)
    
    def preprocess_code(self, code):
        """Clean and preprocess the Python code."""
        # Remove empty lines
        code = re.sub(r'\n\s*\n', '\n', code)
        
        # Remove excessive comments (keeping docstrings)
        code_lines = []
        in_docstring = False
        docstring_delimiter = None
        
        for line in code.split('\n'):
            # Check for docstring delimiters
            if '"""' in line or "'''" in line:
                delimiter = '"""' if '"""' in line else "'''"
                if not in_docstring:
                    in_docstring = True
                    docstring_delimiter = delimiter
                elif docstring_delimiter == delimiter:
                    in_docstring = False
                    docstring_delimiter = None
            
            # Keep docstrings and non-comment lines
            if in_docstring or not line.strip().startswith('#'):
                code_lines.append(line)
        
        processed_code = '\n'.join(code_lines)
        
        # Normalize whitespace
        processed_code = re.sub(r' +', ' ', processed_code)
        
        return processed_code
    
    def extract_functions(self, code):
        """Extract individual functions for summarization"""
        # Simple regex to find function definitions
        function_pattern = r'def\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*\(.*?\).*?:'
        function_matches = re.finditer(function_pattern, code, re.DOTALL)
        
        functions = []
        for match in function_matches:
            start_pos = match.start()
            # Find the function body
            function_name = match.group(1)
            lines = code[start_pos:].split('\n')
            
            # Skip the function definition line
            body_start = 1
            while body_start < len(lines) and not lines[body_start].strip():
                body_start += 1
                
            if body_start < len(lines):
                # Get the indentation of the function body
                body_indent = len(lines[body_start]) - len(lines[body_start].lstrip())
                
                # Gather all lines with at least this indentation
                function_body = [lines[0]]  # The function definition
                i = 1
                while i < len(lines):
                    line = lines[i]
                    if line.strip() and (len(line) - len(line.lstrip())) < body_indent and not line.strip().startswith('#'):
                        break
                    function_body.append(line)
                    i += 1
                
                function_code = '\n'.join(function_body)
                functions.append((function_name, function_code))
        
        # Simple regex to find class methods
        class_pattern = r'class\s+([a-zA-Z_][a-zA-Z0-9_]*)'
        class_matches = re.finditer(class_pattern, code, re.DOTALL)
        
        for match in class_matches:
            class_name = match.group(1)
            start_pos = match.start()
            
            # Find class methods using the function pattern
            class_code = code[start_pos:]
            method_matches = re.finditer(function_pattern, class_code, re.DOTALL)
            
            for method_match in method_matches:
                method_name = method_match.group(1)
                # Skip if this is not a method (i.e., it's a function outside the class)
                if method_match.start() > 200:  # Simple heuristic to check if method is within class scope
                    break
                
                # Get the full method code
                method_start = method_match.start()
                method_lines = class_code[method_start:].split('\n')
                
                # Skip the method definition line
                body_start = 1
                while body_start < len(method_lines) and not method_lines[body_start].strip():
                    body_start += 1
                
                if body_start < len(method_lines):
                    # Get the indentation of the method body
                    body_indent = len(method_lines[body_start]) - len(method_lines[body_start].lstrip())
                    
                    # Gather all lines with at least this indentation
                    method_body = [method_lines[0]]  # The method definition
                    i = 1
                    while i < len(method_lines):
                        line = method_lines[i]
                        if line.strip() and (len(line) - len(line.lstrip())) < body_indent and not line.strip().startswith('#'):
                            break
                        method_body.append(line)
                        i += 1
                    
                    method_code = '\n'.join(method_body)
                    functions.append((f"{class_name}.{method_name}", method_code))
        
        return functions
    
    def extract_classes(self, code):
        """Extract class definitions for summarization"""
        class_pattern = r'class\s+([a-zA-Z_][a-zA-Z0-9_]*)'
        class_matches = re.finditer(class_pattern, code, re.DOTALL)
        
        classes = []
        for match in class_matches:
            class_name = match.group(1)
            start_pos = match.start()
            
            # Extract class body
            class_lines = code[start_pos:].split('\n')
            
            # Skip the class definition line
            body_start = 1
            while body_start < len(class_lines) and not class_lines[body_start].strip():
                body_start += 1
            
            if body_start < len(class_lines):
                # Get the indentation of the class body
                body_indent = len(class_lines[body_start]) - len(class_lines[body_start].lstrip())
                
                # Gather all lines with at least this indentation
                class_body = [class_lines[0]]  # The class definition
                i = 1
                while i < len(class_lines):
                    line = class_lines[i]
                    if line.strip() and (len(line) - len(line.lstrip())) < body_indent:
                        break
                    class_body.append(line)
                    i += 1
                
                class_code = '\n'.join(class_body)
                classes.append((class_name, class_code))
        
        return classes
    
    def summarize(self, code, max_length=50):
        """Generate summary using CodeT5."""
        # Truncate input if needed
        max_input_length = 512  # CodeT5 typically accepts up to 512 tokens
        tokenized_code = self.tokenizer(code, truncation=True, max_length=max_input_length, return_tensors="pt").to(self.device)
        
        with torch.no_grad():
            generated_ids = self.model.generate(
                tokenized_code["input_ids"],
                max_length=max_length,
                num_beams=4,
                early_stopping=True
            )
        
        summary = self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)
        return summary
    
    def summarize_code(self, code, summarize_functions=True, summarize_classes=True):
        """
        Generate full file summary and optionally function/class level summaries.
        Returns a dictionary with summaries.
        """
        preprocessed_code = self.preprocess_code(code)
        
        results = {
            "file_summary": None,
            "function_summaries": {},
            "class_summaries": {}
        }
        
        # Generate file-level summary
        try:
            file_summary = self.summarize(preprocessed_code)
            results["file_summary"] = file_summary
        except Exception as e:
            results["file_summary"] = f"Error generating file summary: {str(e)}"
        
        # Generate function-level summaries if requested
        if summarize_functions:
            functions = self.extract_functions(preprocessed_code)
            
            for function_name, function_code in functions:
                try:
                    summary = self.summarize(function_code)
                    results["function_summaries"][function_name] = summary
                except Exception as e:
                    results["function_summaries"][function_name] = f"Error: {str(e)}"
        
        # Generate class-level summaries if requested
        if summarize_classes:
            classes = self.extract_classes(preprocessed_code)
            
            for class_name, class_code in classes:
                try:
                    summary = self.summarize(class_code)
                    results["class_summaries"][class_name] = summary
                except Exception as e:
                    results["class_summaries"][class_name] = f"Error: {str(e)}"
        
        return results

def main():
    st.set_page_config(
        page_title="Python Code Summarizer",
        page_icon="📝",
        layout="wide"
    )
    
    st.title("📝 Python Code Summarizer using CodeT5")
    st.markdown("""
    Upload a Python file or paste code directly to generate summaries.
    This app uses CodeT5, a pretrained model for code understanding and generation.
    """)
    
    # Initialize session state
    if 'summarizer' not in st.session_state:
        st.session_state.summarizer = None
    
    # Load model if not already loaded
    if st.session_state.summarizer is None:
        st.session_state.summarizer = CodeT5Summarizer()
    
    # Create tabs for different input methods
    tab1, tab2 = st.tabs(["Upload Python File", "Paste Code"])
    
    with tab1:
        uploaded_file = st.file_uploader("Choose a Python file", type=['py'])
        if uploaded_file is not None:
            code = uploaded_file.getvalue().decode('utf-8')
            with st.expander("View Uploaded Code", expanded=False):
                st.code(code, language='python')
            
            # Add summarization options
            st.subheader("Summarization Options")
            col1, col2 = st.columns(2)
            with col1:
                summarize_functions = st.checkbox("Generate function summaries", value=True)
            with col2:
                summarize_classes = st.checkbox("Generate class summaries", value=True)
            
            if st.button("Summarize Code", key="summarize_file"):
                with st.spinner("Generating summaries..."):
                    start_time = time.time()
                    summaries = st.session_state.summarizer.summarize_code(
                        code, 
                        summarize_functions=summarize_functions,
                        summarize_classes=summarize_classes
                    )
                    end_time = time.time()
                    
                    # Display summaries
                    st.success(f"Summarization completed in {end_time - start_time:.2f} seconds!")
                    
                    # File summary
                    st.subheader("File Summary")
                    st.write(summaries["file_summary"])
                    
                    # Function summaries
                    if summarize_functions and summaries["function_summaries"]:
                        st.subheader("Function Summaries")
                        for func_name, summary in summaries["function_summaries"].items():
                            with st.expander(f"Function: {func_name}"):
                                st.write(summary)
                    
                    # Class summaries
                    if summarize_classes and summaries["class_summaries"]:
                        st.subheader("Class Summaries")
                        for class_name, summary in summaries["class_summaries"].items():
                            with st.expander(f"Class: {class_name}"):
                                st.write(summary)
    
    with tab2:
        code = st.text_area("Paste Python code here", height=300)
        if code:
            # Add summarization options
            st.subheader("Summarization Options")
            col1, col2 = st.columns(2)
            with col1:
                summarize_functions = st.checkbox("Generate function summaries", value=True, key="func_paste")
            with col2:
                summarize_classes = st.checkbox("Generate class summaries", value=True, key="class_paste")
            
            if st.button("Summarize Code", key="summarize_paste"):
                with st.spinner("Generating summaries..."):
                    start_time = time.time()
                    summaries = st.session_state.summarizer.summarize_code(
                        code, 
                        summarize_functions=summarize_functions,
                        summarize_classes=summarize_classes
                    )
                    end_time = time.time()
                    
                    # Display summaries
                    st.success(f"Summarization completed in {end_time - start_time:.2f} seconds!")
                    
                    # File summary
                    st.subheader("File Summary")
                    st.write(summaries["file_summary"])
                    
                    # Function summaries
                    if summarize_functions and summaries["function_summaries"]:
                        st.subheader("Function Summaries")
                        for func_name, summary in summaries["function_summaries"].items():
                            with st.expander(f"Function: {func_name}"):
                                st.write(summary)
                    
                    # Class summaries
                    if summarize_classes and summaries["class_summaries"]:
                        st.subheader("Class Summaries")
                        for class_name, summary in summaries["class_summaries"].items():
                            with st.expander(f"Class: {class_name}"):
                                st.write(summary)
    
    st.markdown("---")
    st.markdown("### About")
    st.markdown("""
    This app uses the CodeT5 model to generate summaries of Python code. The model is trained on a large corpus of code and documentation.
    
    **Features:**
    - File-level summaries
    - Function-level summaries
    - Class-level summaries
    
    **Limitations:**
    - Summaries may not always be accurate
    - Long files may be truncated
    - Complex code structures might not be properly understood
    """)

if __name__ == "__main__":
    main()