TharvinPrakash/Multi-model-Chatbot

1import streamlit as st2from transformers import AutoTokenizer, TFAutoModelForSeq2SeqLM, AutoModelForSeq2SeqLM3import torch4 5# Load Hugging Face tokenizer and model for re-punctuation6@st.cache_resource7def load_re_punctuate_model():8    tokenizer = AutoTokenizer.from_pretrained("SJ-Ray/Re-Punctuate")9    model = TFAutoModelForSeq2SeqLM.from_pretrained("SJ-Ray/Re-Punctuate")10    return tokenizer, model11 12# Load Hugging Face tokenizer and model for headline generation (local path)13@st.cache_resource14def load_headline_model(model_path):15    tokenizer = AutoTokenizer.from_pretrained(model_path)16    model = AutoModelForSeq2SeqLM.from_pretrained(model_path)17    return tokenizer, model18 19# Function to re-punctuate text20def re_punctuate_text(tokenizer, model, text):21    inputs = tokenizer(text, return_tensors="tf", max_length=512, truncation=True)22    outputs = model.generate(inputs["input_ids"], max_length=512, num_beams=4, early_stopping=True)23    return tokenizer.decode(outputs[0], skip_special_tokens=True)24 25# Function to generate a headline26def generate_headline_text(tokenizer, model, text, max_length=50):27    inputs = tokenizer(f"headline: {text}", return_tensors="pt", truncation=True, padding=True)28    with torch.no_grad():29        outputs = model.generate(30             **inputs,31            max_length=max_length,32            num_beams=5,33            no_repeat_ngram_size=2,34            early_stopping=True35        )36    return tokenizer.decode(outputs[0], skip_special_tokens=True)37 38# Streamlit app layout39st.title("Model Selection: Re-Punctuate or Generate Headline")40 41# Model selection dropdown42model_options = ["Re-Punctuate Text", "Generate Headline"]43selected_model = st.selectbox("Choose a model to use:", model_options)44 45# User input text46input_text = st.text_area("Enter text:", placeholder="Type your input here...")47 48# Default local model path for headline generation49local_model_path = "Michau/t5-base-en-generate-headline"50 51# Button to process text based on the selected model52if st.button("Process Text") and input_text:53    with st.spinner("Processing..."):54        if selected_model == "Re-Punctuate Text":55            tokenizer, model = load_re_punctuate_model()56            result = re_punctuate_text(tokenizer, model, input_text)57        else:  # Generate Headline58            tokenizer, model = load_headline_model(local_model_path)59            result = generate_headline_text(tokenizer, model, input_text)60        61        # Display result62        st.subheader(f"Result from {selected_model}:")63        st.write(result)64 65# Footer66st.write("---")67st.write("Powered by Hugging Face Models.")68