app.py

import streamlit as st
from st_utils import (
    load_tokenizer_and_model,
    generate_docstring,
    download_model,
    # list_files,
)
from huggingface_hub import hf_hub_download
import os
import torch

# list_files(os.getcwd())

# Set the title and description of the app
st.title("Code Function Summarization App")
st.write(
    """
This app uses the Hugging Face transformers library to generate summaries of input text. 
Simply select one of the sample Python functions from the dropdown menu below, and click the 'Summarize' button to generate a summary for the corresponding function.
"""
)

# st.write(f"Has CUDA: {torch.cuda.is_available()}")

# Download the model from the Hugging Face Hub if it doesn't exist
download_model()

# load the tokenizer and model
tokenizer, model, device = load_tokenizer_and_model("./models/pytorch_model.bin")

# Create a dropdown menu for the user to select a sample Python function
values = [
    "",
    "def multiply(a, b):\n    return a * b",
    "def get_data():\n    data = []\n    for i in range(10):\n        data.append(i)\n    return data",
    "def search(data, target):\n    for i in range(len(data)):\n        if data[i] == target:\n            return i\n    return -1",
]

selected_value = st.selectbox("Select a sample Python function:", values)

# Create a text input area for the user to enter their text
text_input = st.text_area(
    "Or enter your Python function here (prioritize this over the dropdown menu):",
    height=256,
    value=selected_value,
)


# Define a function to generate a summary
def generate_summary(text):
    summary = generate_docstring(model, tokenizer, device, text, max_length=30)
    return summary


# When the user clicks the 'Summarize' button, generate a summary
if st.button("Summarize") and (len(selected_value) > 0 or len(text_input) > 0):
    with st.spinner("Generating summary..."):
        if len(text_input) > 0:
            summaries = generate_summary(text_input)
            st.subheader("Docstrings:")
            for i, summary in enumerate(summaries):
                st.write(f"{i + 1}. " + summary)
        # if len(selected_value) > 0:
        else:
            summaries = generate_summary(selected_value)
            st.subheader("Docstrings:")
            for i, summary in enumerate(summaries):
                st.write(f"{i + 1}. " + summary)