Update app.py

app.py

@@ -1,17 +1,14 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
 # 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")
+tokenizer = AutoTokenizer.from_pretrained("sagorsarker/emailgenerator")
+model = AutoModelForCausalLM.from_pretrained("sagorsarker/emailgenerator")
 
 # Streamlit UI styling
-st.set_page_config(page_title="ReplyCraft Email Generator", layout="centered")
+st.set_page_config(page_title="Email Generator", layout="centered")
 
-# Add custom CSS for styling
+# Add some custom CSS for styling
 st.markdown("""
     <style>
         .title {
@@ -43,28 +40,30 @@ st.markdown("""
 """, unsafe_allow_html=True)
 
 # Title and Header
-st.markdown('<div class="title">ReplyCraft - Email Reply Generator</div>', unsafe_allow_html=True)
+st.markdown('<div class="title">Email 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 Reply", key="generate_button", help="Click to generate reply", use_container_width=True)
+generate_button = st.button("Generate Email", key="generate_button", help="Click to generate email", use_container_width=True)
 
 # Handling the generation
 if generate_button:
     if user_input:
-        # Generate the reply email
-        result = generator(
-            user_input,
-            max_length=150,
-            do_sample=True,
-            early_stopping=True,
-        )
+        # 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)
         
-        # 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)
+        # 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)
     else:
-        st.error("Please enter a prompt to generate the email reply.")
+        st.error("Please enter a prompt to generate the email.")
+