Update app.py

app.py

@@ -1,24 +1,24 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+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 = AutoModelForSeq2SeqLM.from_pretrained("SJ-Ray/Re-Punctuate",from_tf=True)
+    model = TFAutoModelForSeq2SeqLM.from_pretrained("SJ-Ray/Re-Punctuate")
     return tokenizer, model
 
-# Load Hugging Face tokenizer and model for headline generation
+# Load Hugging Face tokenizer and model for headline generation (local path)
 @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")
+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="pt", max_length=512, truncation=True)
+    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)
 
@@ -27,7 +27,7 @@ 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,
+             **inputs,
             max_length=max_length,
             num_beams=5,
             no_repeat_ngram_size=2,
@@ -45,6 +45,9 @@ 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 = r"C:\Users\Tharvin prakash\.cache\huggingface\hub\models--Michau--t5-base-en-generate-headline\snapshots\f526532f788c45b6b6288286e5ef929fa768ef6a"
+
 # Button to process text based on the selected model
 if st.button("Process Text") and input_text:
     with st.spinner("Processing..."):
@@ -52,7 +55,7 @@ if st.button("Process Text") and input_text:
             tokenizer, model = load_re_punctuate_model()
             result = re_punctuate_text(tokenizer, model, input_text)
         else:  # Generate Headline
-            tokenizer, model = load_headline_model()
+            tokenizer, model = load_headline_model(local_model_path)
             result = generate_headline_text(tokenizer, model, input_text)
         
         # Display result