Update app.py

app.py

@@ -1,97 +1,50 @@
 import streamlit as st
-from huggingface_hub import InferenceClient
-import os
+from transformers import pipeline
 
-# Get Hugging Face API token from Space secrets
-HF_TOKEN = os.getenv("HF_TOKEN")
-client = InferenceClient(token=HF_TOKEN)
+# Cache the model to load it only once
+@st.cache_resource
+def load_generator():
+    return pipeline("text-generation", model="aubmindlab/aragpt2-base", device=-1)  # device=-1 forces CPU
 
-# Initialize session state
-if "chat_history" not in st.session_state:
-    st.session_state.chat_history = []
-if "corrected_sentence" not in st.session_state:
-    st.session_state.corrected_sentence = ""
+# Load the text generation pipeline
+generator = load_generator()
 
-# Title of the app
-st.title("Sentence Improver & Chat with DeepSeek (Hugging Face)")
-
-# --- Sentence Correction Section ---
-st.subheader("Improve a Sentence")
-user_input = st.text_input("Enter a sentence to improve:", "I goed to the park and play.")
+# App title
+st.title("Arabic Sentence Improver & Chat App")
 
+# Sentence Correction Section
+st.subheader("Improve an Arabic Sentence")
+user_input = st.text_input("Enter an Arabic sentence to improve:", "أنا ذهبت الحديقة")
 if st.button("Improve Sentence"):
     if user_input:
-        prompt = f"Correct and improve this sentence: '{user_input}'"
+        # Prompt the model to correct the sentence
+        prompt = f"Correct this Arabic sentence: '{user_input}' to"
         try:
-            response = client.text_generation(
-                prompt,
-                model="aubmindlab/aragpt2-base",  # DeepSeek model on Hugging Face
-                max_new_tokens=100,
-                temperature=0.7,
-            )
-            st.session_state.corrected_sentence = response.strip()
-            st.success(f"Improved Sentence: {st.session_state.corrected_sentence}")
+            response = generator(prompt, max_new_tokens=50, temperature=0.7)[0]["generated_text"]
+            # Extract the corrected sentence (assuming it follows "to")
+            corrected_sentence = response.split("to")[1].strip() if "to" in response else response
+            st.session_state.corrected_sentence = corrected_sentence
+            st.success(f"Improved Sentence: {corrected_sentence}")
         except Exception as e:
             st.error(f"Error: {str(e)}")
     else:
         st.warning("Please enter a sentence first!")
 
-# --- Chat Section ---
+# Chat Section
 st.subheader("Chat About the Corrected Sentence")
-if st.session_state.corrected_sentence:
-    # Chat history container with scrollbar
-    chat_container = st.container(height=300)
-    with chat_container:
-        for speaker, message in st.session_state.chat_history:
-            if speaker == "You":
-                st.markdown(
-                    f"<div style='text-align: right; margin: 5px;'><span style='background-color: #DCF8C6; padding: 8px; border-radius: 10px;'>{message}</span></div>",
-                    unsafe_allow_html=True
-                )
-            else:
-                st.markdown(
-                    f"<div style='text-align: left; margin: 5px;'><span style='background-color: #ECECEC; padding: 8px; border-radius: 10px;'>{message}</span></div>",
-                    unsafe_allow_html=True
-                )
-
-    # Chat input with Enter submission
-    chat_input = st.text_input(
-        "Ask something about the corrected sentence (press Enter to send) ➡️",
-        key="chat_input",
-        value="",
-        on_change=lambda: submit_chat(),
-    )
-
-    # Function to handle chat submission with full context
-    def submit_chat():
-        chat_text = st.session_state.chat_input
-        if chat_text:
-            # Build prompt with corrected sentence and full chat history
-            prompt = f"The corrected sentence is: '{st.session_state.correctedSentence}'.\n"
-            if st.session_state.chat_history:
-                prompt += "Previous conversation:\n"
-                for speaker, message in st.session_state.chat_history:
-                    prompt += f"{speaker}: {message}\n"
-            prompt += f"User asks: '{chat_text}'. Respond naturally."
-            
+if "corrected_sentence" in st.session_state:
+    chat_input = st.text_input("Ask something about the corrected sentence:", key="chat_input")
+    if st.button("Send"):
+        if chat_input:
+            # Prompt the model with context for chatting
+            prompt = f"The corrected sentence is: '{st.session_state.corrected_sentence}'. User asks: '{chat_input}'"
             try:
-                response = client.text_generation(
-                    prompt,
-                    model="mistralai/Mixtral-8x7B-Instruct-v0.1",
-                    max_new_tokens=150,
-                    temperature=0.7,
-                )
-                llm_response = response.strip()
-                st.session_state.chat_history.append(("You", chat_text))
-                st.session_state.chat_history.append(("LLM", llm_response))
-                st.session_state.chat_input = ""
+                response = generator(prompt, max_new_tokens=100, temperature=0.7)[0]["generated_text"]
+                st.write(f"**You:** {chat_input}")
+                st.write(f"**LLM:** {response}")
             except Exception as e:
                 st.error(f"Error in chat: {str(e)}")
-
+        else:
+            st.warning("Please enter a question!")
 else:
-    st.write("Improve a sentence first to start chatting!")
-
-# Optional: Clear chat button
-if st.button("Clear Chat"):
-    st.session_state.chat_history = []
-    st.rerun()
+    st.write("Improve a sentence first to start chatting!")