Spaces:
Sleeping
Sleeping
File size: 2,629 Bytes
adce22c 5020ea6 adce22c 5020ea6 adce22c 5020ea6 adce22c 5020ea6 adce22c 5020ea6 adce22c 5020ea6 adce22c 5020ea6 ef6455e 5020ea6 adce22c 5020ea6 adce22c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import streamlit as st
from transformers import AutoTokenizer, TFAutoModelForSeq2SeqLM, AutoModelForSeq2SeqLM
import torch
# Load Hugging Face tokenizer and model for re-punctuation
@st.cache_resource
def load_re_punctuate_model():
tokenizer = AutoTokenizer.from_pretrained("SJ-Ray/Re-Punctuate")
model = TFAutoModelForSeq2SeqLM.from_pretrained("SJ-Ray/Re-Punctuate")
return tokenizer, model
# Load Hugging Face tokenizer and model for headline generation (local path)
@st.cache_resource
def load_headline_model(model_path):
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
return tokenizer, model
# Function to re-punctuate text
def re_punctuate_text(tokenizer, model, text):
inputs = tokenizer(text, return_tensors="tf", max_length=512, truncation=True)
outputs = model.generate(inputs["input_ids"], max_length=512, num_beams=4, early_stopping=True)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# Function to generate a headline
def generate_headline_text(tokenizer, model, text, max_length=50):
inputs = tokenizer(f"headline: {text}", return_tensors="pt", truncation=True, padding=True)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_length=max_length,
num_beams=5,
no_repeat_ngram_size=2,
early_stopping=True
)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# Streamlit app layout
st.title("Model Selection: Re-Punctuate or Generate Headline")
# Model selection dropdown
model_options = ["Re-Punctuate Text", "Generate Headline"]
selected_model = st.selectbox("Choose a model to use:", model_options)
# User input text
input_text = st.text_area("Enter text:", placeholder="Type your input here...")
# Default local model path for headline generation
local_model_path = "Michau/t5-base-en-generate-headline"
# Button to process text based on the selected model
if st.button("Process Text") and input_text:
with st.spinner("Processing..."):
if selected_model == "Re-Punctuate Text":
tokenizer, model = load_re_punctuate_model()
result = re_punctuate_text(tokenizer, model, input_text)
else: # Generate Headline
tokenizer, model = load_headline_model(local_model_path)
result = generate_headline_text(tokenizer, model, input_text)
# Display result
st.subheader(f"Result from {selected_model}:")
st.write(result)
# Footer
st.write("---")
st.write("Powered by Hugging Face Models.")
|