app.py

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.")