Update app.py

app.py

@@ -1,33 +1,91 @@
 import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-# Load the merged model
-model_name = "EmTpro01/gemma-paraphraser-4bit"  # Replace with your merged model path
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)  # Default device is CPU
-
-# Streamlit UI
-st.title("Text Paraphrasing ")
-st.write("Provide a paragraph, and this AI will paraphrase it for you.")
-
-# Input paragraph
-paragraph = st.text_area("Enter a paragraph to paraphrase:", height=200)
-
-if st.button("Paraphrase"):
-    if paragraph.strip():
-        with st.spinner("Paraphrasing..."):
-            # Prepare the prompt
-            alpaca_prompt = f"Below is a paragraph, paraphrase it.\n### paragraph: {paragraph}\n### paraphrased:"
-            
-            # Tokenize input and move to CPU
-            inputs = tokenizer(alpaca_prompt, return_tensors="pt")
-            
-            # Generate paraphrased text
-            output = model.generate(**inputs, max_new_tokens=200)
-            paraphrased = tokenizer.decode(output[0], skip_special_tokens=True)
-            
-            # Extract the paraphrased portion
-            result = paraphrased.split("### paraphrased:")[1].strip()
-            st.text_area("Paraphrased Output:", result, height=200)
-    else:
-        st.warning("Please enter a paragraph to paraphrase.")
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+# Set page configuration
+st.set_page_config(page_title="Gemma Paraphraser", page_icon="✍️")
+
+# Load model and tokenizer
+@st.cache_resource
+def load_model():
+    model_name = "EmTpro01/gemma-paraphraser-16bit"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name, 
+        device_map="cpu",
+        torch_dtype=torch.float16
+    )
+    return model, tokenizer
+
+# Paraphrase function
+def paraphrase_text(text, model, tokenizer):
+    # Prepare the prompt using Alpaca format
+    system_prompt = "Below is provided a paragraph, paraphrase it"
+    prompt = f"{system_prompt}\n\n### Input:\n{text}\n\n### Output:\n"
+    
+    # Tokenize input
+    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
+    
+    # Generate paraphrased text
+    outputs = model.generate(
+        inputs.input_ids, 
+        max_length=512,  # Adjust based on your needs
+        num_return_sequences=1,
+        temperature=0.7,
+        do_sample=True
+    )
+    
+    # Decode and clean the output
+    paraphrased = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract the output part (after "### Output:")
+    output_start = paraphrased.find("### Output:") + len("### Output:")
+    paraphrased_text = paraphrased[output_start:].strip()
+    
+    return paraphrased_text
+
+# Streamlit App
+def main():
+    st.title("📝 Gemma Paraphraser")
+    st.write("Paraphrase your text using the Gemma model")
+
+    # Load model
+    try:
+        model, tokenizer = load_model()
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
+        return
+
+    # Input text area
+    input_text = st.text_area("Enter text to paraphrase:", height=200)
+
+    # Paraphrase button
+    if st.button("Paraphrase"):
+        if input_text:
+            with st.spinner("Generating paraphrase..."):
+                try:
+                    paraphrased_text = paraphrase_text(input_text, model, tokenizer)
+                    
+                    # Display results
+                    st.subheader("Paraphrased Text:")
+                    st.write(paraphrased_text)
+                    
+                    # Optional: Copy to clipboard
+                    st.button("Copy to Clipboard", 
+                              on_click=lambda: st.write(paraphrased_text))
+                except Exception as e:
+                    st.error(f"Error during paraphrasing: {e}")
+        else:
+            st.warning("Please enter some text to paraphrase.")
+
+    # Additional information
+    st.sidebar.info(
+        "Model: EmTpro01/gemma-paraphraser-16bit\n\n"
+        "Tips:\n"
+        "- Enter a paragraph to paraphrase\n"
+        "- Click 'Paraphrase' to generate\n"
+        "- Running on CPU with 16-bit precision"
+    )
+
+if __name__ == "__main__":
+    main()