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

# list_files(os.getcwd())

# Set the title and description of the app
st.title("Text 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.
"""
)

# 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",
]

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

# Create a text input area for the user to enter their text
text_input = st.text_area(
    "Or enter your Python function here:",
    height=300,
    value=values[0],
)


# 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(selected_value) > 0:
            summaries = generate_summary(selected_value)
            st.subheader("Docstrings:")
            for i, summary in enumerate(summaries):
                st.write(f"{i + 1}. " + summary)
        else:
            summaries = generate_summary(text_input)
            st.subheader("Docstrings:")
            for i, summary in enumerate(summaries):
                st.write(f"{i + 1}. " + summary)


# import streamlit as st
# from st_utils import load_tokenizer_and_model, generate_docstring, download_model

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


# # Set the title and description of the app
# st.title("Text Summarization App")
# st.write(
#     """
# This app uses the Hugging Face transformers library to generate summaries of input text.
# Simply enter your text in the input area below, and click the 'Summarize' button to generate a summary.
# """
# )

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

# # Create a text input area for the user to enter their text
# 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",
# ]

# st.subheader("Enter your Python function here:")
# text_input = st.text_area(
#     "Input text here...",
#     height=300,
#     value=values[2],
# )


# # 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(text_input) > 0:
#     with st.spinner("Generating summary..."):
#         # summary = generate_summary(text_input)
#         # st.write("Summary:")
#         # st.code(summary, language="text")
#         summaries = generate_summary(text_input)
#         st.subheader("Summary:")
#         for i, summary in enumerate(summaries):
#             st.write(f"{i + 1}. " + summary)