Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import T5Tokenizer, TFAutoModelForSeq2SeqLM, pipeline | |
import zipfile | |
import os | |
import nltk | |
# Download NLTK data | |
nltk.download('punkt') | |
from nltk.tokenize import sent_tokenize | |
# Define the path to the saved model zip file (ensure there is no extra space) | |
zip_model_path = 'bartzip.zip' | |
# Define the directory to extract the model | |
model_dir = './model' | |
# Unzip the model | |
with zipfile.ZipFile(zip_model_path, 'r') as zip_ref: | |
zip_ref.extractall(model_dir) | |
# After unzipping, the model should be in a specific directory, check the directory structure | |
model_path = os.path.join(model_dir, 'Bart_model') | |
# Print out the model_path for debugging | |
print("Model Path:", model_path) | |
# Verify that the directory exists and contains the necessary files | |
if not os.path.exists(model_path): | |
st.error(f"Model directory {model_path} does not exist or is incorrect.") | |
# Print out contents of model_dir for further debugging | |
# print("Contents of model_dir:", os.listdir(model_dir)) | |
else: | |
# Load the tokenizer and model from the extracted directory | |
tokenizer = T5Tokenizer.from_pretrained(model_path) | |
model = TFAutoModelForSeq2SeqLM.from_pretrained(model_path) | |
# Create a summarization pipeline | |
summarizer = pipeline("summarization", model=model, tokenizer=tokenizer) | |
# Set the title for the Streamlit app | |
st.title("BART Summary Generator") | |
# Text input for the user | |
text = st.text_area("Enter your text: ") | |
def generate_summary(input_text): | |
# Perform summarization | |
summary = summarizer(input_text, max_length=200, min_length=40, do_sample=False) | |
return summary[0]['summary_text'] | |
if st.button("Generate"): | |
if text: | |
generated_summary = generate_summary(text) | |
# Display the generated summary | |
st.subheader("Generated Summary") | |
st.write(generated_summary) | |
else: | |
st.warning("Please enter some text to generate a summary.") | |