Create app.py

app.py

@@ -0,0 +1,64 @@
+import streamlit as st
+from transformers import AutoTokenizer, 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 = AutoModelForSeq2SeqLM.from_pretrained("SJ-Ray/Re-Punctuate")
+    return tokenizer, model
+
+# Load Hugging Face tokenizer and model for headline generation
+@st.cache_resource
+def load_headline_model():
+    tokenizer = AutoTokenizer.from_pretrained("Michau/t5-base-en-generate-headline")
+    model = AutoModelForSeq2SeqLM.from_pretrained("Michau/t5-base-en-generate-headline")
+    return tokenizer, model
+
+# Function to re-punctuate text
+def re_punctuate_text(tokenizer, model, text):
+    inputs = tokenizer(text, return_tensors="pt", 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...")
+
+# 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()
+            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.")