Update app.py

app.py

@@ -1,14 +1,17 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 # Load the model and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("postbot/gpt2-medium-emailgen")
 model = AutoModelForCausalLM.from_pretrained("postbot/gpt2-medium-emailgen")
 
+# Initialize the text generation pipeline
+generator = pipeline('text-generation', model="postbot/gpt2-medium-emailgen")
+
 # Streamlit UI styling
-st.set_page_config(page_title="Email Generator", layout="centered")
+st.set_page_config(page_title="ReplyCraft Email Generator", layout="centered")
 
-# Add some custom CSS for styling
+# Add custom CSS for styling
 st.markdown("""
     <style>
         .title {
@@ -40,28 +43,28 @@ st.markdown("""
 """, unsafe_allow_html=True)
 
 # Title and Header
-st.markdown('<div class="title">Email Generator</div>', unsafe_allow_html=True)
+st.markdown('<div class="title">ReplyCraft - Email Reply Generator</div>', unsafe_allow_html=True)
 
 # User input for email prompt
 user_input = st.text_area("Enter the email content prompt:", height=150, key="email_prompt", max_chars=500)
 
 # Add a styled button
-generate_button = st.button("Generate Email", key="generate_button", help="Click to generate email", use_container_width=True)
+generate_button = st.button("Generate Email Reply", key="generate_button", help="Click to generate reply", use_container_width=True)
 
 # Handling the generation
 if generate_button:
     if user_input:
-        # Tokenize the input
-        inputs = tokenizer.encode(user_input, return_tensors="pt")
-        
-        # Generate the email text
-        outputs = model.generate(inputs, max_length=300, num_return_sequences=1, no_repeat_ngram_size=2)
+        # Generate the reply email
+        result = generator(
+            user_input,
+            max_length=150,
+            do_sample=True,
+            early_stopping=True,
+        )
         
-        # Decode and display the result
-        generated_email = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Display the generated email with some styling
-        st.markdown('<div class="header">Generated Email:</div>', unsafe_allow_html=True)
-        st.markdown(f'<div>{generated_email}</div>', unsafe_allow_html=True)
+        # Display the generated email reply with styling
+        st.markdown('<div class="header">Generated Email Reply:</div>', unsafe_allow_html=True)
+        st.markdown(f'<div">{result[0]["generated_text"]}</div>', unsafe_allow_html=True)
     else:
-        st.error("Please enter a prompt to generate the email.")
+        st.error("Please enter a prompt to generate the email reply.")
+