Update app.py

app.py

@@ -1,73 +1,140 @@
 import streamlit as st
-from streamlit_chat import message
-from streamlit_extras.colored_header import colored_header
-from streamlit_extras.add_vertical_space import add_vertical_space
-from hugchat import hugchat
-from hugchat.login import Login
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from datasets import load_dataset
+from gtts import gTTS
+import os
+import re
+import random
 
-st.set_page_config(page_title="HugChat - An LLM-powered Streamlit app")
+# Enable Dark Mode and Custom CSS
+st.markdown(
+    """
+    <style>
+        body {
+            background-color: #121212;
+            color: white;
+        }
+        .css-1d391kg {
+            background-color: #333;
+        }
+        .stButton > button {
+            background-color: #6200ee;
+            color: white;
+        }
+        .stTextInput input {
+            background-color: #333;
+            color: white;
+        }
+    </style>
+    """,
+    unsafe_allow_html=True,
+)
 
-# Sidebar contents
-with st.sidebar:
-    st.title('🤗💬 HugChat App')
-    
-    st.header('Hugging Face Login')
-    hf_email = st.text_input('Enter E-mail:', type='password')
-    hf_pass = st.text_input('Enter password:', type='password')
-    
-    st.markdown('''
-    ## About
-    This app is an LLM-powered chatbot built using:
-    - [Streamlit](https://streamlit.io/)
-    - [HugChat](https://github.com/Soulter/hugging-chat-api)
-    - [OpenAssistant/oasst-sft-6-llama-30b-xor](https://huggingface.co/OpenAssistant/oasst-sft-6-llama-30b-xor) LLM model
+# Load models and datasets
+try:
+    code_llama_model = AutoModelForCausalLM.from_pretrained("meta-llama/CodeLlama-7B-Python")
+    code_llama_tokenizer = AutoTokenizer.from_pretrained("meta-llama/CodeLlama-7B-Python")
+except Exception as e:
+    st.error(f"Error loading model: {e}")
+    code_llama_model = None
+    code_llama_tokenizer = None
 
-    ''')
-    add_vertical_space(5)
-    st.write('Made with ❤️ by [Data Professor](https://youtube.com/dataprofessor)')
+try:
+    wordlist_dataset = load_dataset("Canstralian/Wordlists")
+except Exception as e:
+    st.error(f"Error loading Wordlist dataset: {e}")
+    wordlist_dataset = None
 
-# Generate empty lists for generated and past.
-## generated stores AI generated responses
-if 'generated' not in st.session_state:
-    st.session_state['generated'] = ["I'm HugChat, How may I help you?"]
-## past stores User's questions
-if 'past' not in st.session_state:
-    st.session_state['past'] = ['Hi!']
+# Initialize chat history storage
+if "messages" not in st.session_state:
+    st.session_state.messages = [{"role": "assistant", "content": "How may I assist you?"}]
 
-# Layout of input/response containers
-input_container = st.container()
-colored_header(label='', description='', color_name='blue-30')
-response_container = st.container()
+# Function to validate the prompt using regular expressions
+def validate_prompt(prompt: str) -> bool:
+    """
+    Validates if the input prompt is not empty and meets some basic format rules.
+    Args:
+        prompt (str): The input prompt to be validated.
+    Returns:
+        bool: True if the prompt is valid, False otherwise.
+    """
+    # Improved validation: Allow alphanumeric characters, spaces, and punctuation
+    if re.match(r'^[A-Za-z0-9\s\.,;!?(){}[\]]+$', prompt):
+        return True
+    return False
 
-# User input
-## Function for taking user provided prompt as input
-def get_text():
-    input_text = st.text_input("You: ", "", key="input")
-    return input_text
-## Applying the user input box
-with input_container:
-    user_input = get_text()
+# Function to convert text to speech
+def text_to_speech(text: str) -> None:
+    """
+    Converts text to speech using gTTS and saves it as an MP3 file.
+    Args:
+        text (str): The text to be converted to speech.
+    """
+    try:
+        tts = gTTS(text, lang='en')
+        tts.save("response.mp3")
+        os.system("mpg321 response.mp3")
+    except Exception as e:
+        st.error(f"Error generating speech: {e}")
+
+# Display chat history
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.write(message["content"])
+
+# Function to generate chatbot response
+def generate_response(prompt: str) -> str:
+    """
+    Generates a response from the assistant based on the user input.
+    Args:
+        prompt (str): The user's input prompt.
+    Returns:
+        str: The generated response from the assistant.
+    """
+    if code_llama_model and code_llama_tokenizer:
+        if "python" in prompt.lower():
+            # Use the Code Llama model for code-related queries
+            inputs = code_llama_tokenizer(prompt, return_tensors="pt")
+            outputs = code_llama_model.generate(**inputs, max_length=150, num_return_sequences=1)
+            response = code_llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        else:
+            response = "I'm here to assist with your queries."
+    else:
+        response = "Model not loaded. Please try again later."
+
+    if "osint" in prompt.lower():
+        # Respond with dataset-based OSINT information
+        response = "OSINT data analysis coming soon!"
+    elif "wordlist" in prompt.lower() and wordlist_dataset:
+        # Fetch and display a random entry from the Wordlist dataset
+        wordlist_entry = random.choice(wordlist_dataset["train"])["text"]
+        response = f"Here's a random wordlist entry: {wordlist_entry}"
 
-# Response output
-## Function for taking user prompt as input followed by producing AI generated responses
-def generate_response(prompt, email, passwd):
-    # Hugging Face Login
-    sign = Login(email, passwd)
-    cookies = sign.login()
-    sign.saveCookies()
-    # Create ChatBot                        
-    chatbot = hugchat.ChatBot(cookies=cookies.get_dict())
-    response = chatbot.chat(prompt)
     return response
 
-## Conditional display of AI generated responses as a function of user provided prompts
-with response_container:
-    if user_input and hf_email and hf_pass:
-        response = generate_response(user_input, hf_email, hf_pass)
-        st.session_state.past.append(user_input)
-        st.session_state.generated.append(response)
+# User input handling
+if prompt := st.chat_input():
+    # Validate user input
+    if validate_prompt(prompt):
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.write(prompt)
+
+        # Generate and display response with smooth animations
+        with st.chat_message("assistant"):
+            with st.spinner("Assistant is typing..."):
+                response = generate_response(prompt)
+                st.write(response)
         
-    if st.session_state['generated']:
-        for i in range(len(st.session_state['generated'])):
-            message(st.session_state['past'][i], is_user=True, key=str(i) + '_user')
-            message(st.session_state["generated"][i], key=str(i))
+        # Text-to-Speech integration for the assistant's response
+        text_to_speech(response)
+
+        # Store the assistant's response
+        st.session_state.messages.append({"role": "assistant", "content": response})
+    else:
+        st.warning("Invalid input. Please ensure your input contains only valid characters.")
+
+# User Feedback Section
+feedback = st.selectbox("How was your experience?", ["😊 Excellent", "😐 Okay", "😕 Poor"])
+if feedback:
+    st.success(f"Thank you for your feedback: {feedback}", icon="✅")