Amrabdellatief/python-code-to-uml

1import streamlit as st2from transformers import pipeline3from graphviz import Digraph4import re5 6# Load the pre-trained CodeBERT model7model_name = "microsoft/codebert-base"8code_parser = pipeline("text2text-generation", model=model_name)9 10def extract_classes_methods(code):11    # Use CodeBERT model to parse code and generate class/method information12    result = code_parser(code)13    parsed_code = result[0]["generated_text"]14    15    return parsed_code16 17def generate_uml(parsed_code):18    # Initialize Graphviz Digraph19    dot = Digraph(format='png')20 21    # Regex patterns to find classes and methods22    class_pattern = re.compile(r'class (\w+):')23    method_pattern = re.compile(r'(\w+)\(.*\):')24 25    # Extract class and method names26    classes = class_pattern.findall(parsed_code)27    methods = method_pattern.findall(parsed_code)28 29    # Add classes and methods to the UML diagram30    for class_name in classes:31        dot.node(class_name, f'class {class_name}')32    33    for method_name in methods:34        dot.node(method_name, f'  + {method_name}()')35        # Link methods to the classes36        for class_name in classes:37            dot.edge(class_name, method_name)38 39    return dot40 41# Streamlit app42def main():43    st.title("Python Code to UML Diagram")44    45    # Input Python code46    code = st.text_area("Enter Python Code", height=300)47 48    if st.button("Generate UML"):49        if code:50            # Extract classes and methods using CodeBERT51            parsed_code = extract_classes_methods(code)52 53            # Generate the UML diagram54            uml_diagram = generate_uml(parsed_code)55 56            # Render the UML diagram to an image and display it57            uml_diagram.render('/kaggle/working/uml_class_diagram')58            st.image('/kaggle/working/uml_class_diagram.png', caption='UML Class Diagram')59        else:60            st.warning("Please enter some Python code.")61 62if __name__ == "__main__":63    main()64