change to gemini

chat-streamlit-app.py

@@ -1,106 +1,64 @@
 import streamlit as st
-import logging
-import time
-from transformers import pipeline
-import re
-from huggingface_hub import login
-from dotenv import load_dotenv
-import os
+import google.generativeai as genai
 
-
-logging.basicConfig(level=logging.INFO)
-#hf_token = os.getenv('HUGGINGFACE_TOKEN')
-
-# If the token is found, it will be used
-#if hf_token:
-#    logging.info("Hugging Face Token retrieved successfully!")
-#else:
-#    logging.info("Error: Hugging Face Token not found!")
-#login(token=st.secrets["HUGGINGFACE_TOKEN"])
-
-#     "Meta-LLaMA 3 8B": "meta-llama/Meta-Llama-3-8B-Instruct",
-MODEL_DICT = {
-    "DialoGPT-medium": "microsoft/DialoGPT-medium",
-    "flan-t5-small": "google/flan-t5-small",
-    "Meta-Llama-Guard-2-8B": "meta-llama/Meta-Llama-Guard-2-8B"
-}
-
-# Initialize chat history
-if 'messages' not in st.session_state:
+# Initialize session state variable for messages
+if "messages" not in st.session_state:
     st.session_state.messages = []
 
-if 'current_model' not in st.session_state:
-    st.session_state.current_model = None
-
-
-def load_generator(model_name):
-    return pipeline("text-generation", model=model_name)
-
-
-def stream_chat(generator, messages, max_context=5, preprompt="You are a helpful and friendly chatbot."):
-    try:
-        history = ""
-        for msg in messages[-max_context:]:
-            if msg['role'] == 'user':
-                history += f"User: {msg['content']}{generator.tokenizer.eos_token}\n"
-            elif msg['role'] == 'assistant':
-                history += f"Assistant: {msg['content']}{generator.tokenizer.eos_token}\n"
-
-        prompt = f"{preprompt} User: {messages[-1]['content']}{generator.tokenizer.eos_token}\nAssistant:"
-
-        output = generator(prompt, max_length=150, do_sample=True, temperature=0.7, pad_token_id=generator.tokenizer.eos_token_id)
-        response = output[0]['generated_text']
-        response = response.split("Assistant:")[-1].strip()
-        response = response.replace(generator.tokenizer.eos_token, "").strip()
-
-        return response
-    except Exception as e:
-        logging.error(f"Error during generation: {str(e)}")
-        raise e
+def create_sidebar():
+    """Creates the sidebar for the app."""
+    with st.sidebar:
+        st.title("🤖 Gemini Chatbot")
+        api_key = st.text_input("Google API Key:", type="password", help="Get your API key from Google Cloud")
+        st.markdown("""
+        ### What is this?
+        A simple chatbot powered by the Gemini language model.
+        
+        ### How to use
+        1.  Enter your Google API key in the sidebar.
+        2.  Start chatting!
+        """)
+        return api_key
 
 def main():
-    st.title("Chat with Multiple LLMs")
-    logging.info("App started")
-
-    # Sidebar for model selection
-    selected_model_key = st.sidebar.selectbox("Choose a model", list(MODEL_DICT.keys()))
-    selected_model_name = MODEL_DICT[selected_model_key]
-
-    # Load generator only if model changed
-    if st.session_state.current_model != selected_model_name:
-        with st.spinner(f"Loading model {selected_model_key}..."):
-            #hf_token = os.environ.get("HUGGINGFACE_TOKEN")
-            st.session_state.generator = load_generator(selected_model_name)
-        st.session_state.current_model = selected_model_name
-        st.success(f"Model {selected_model_key} loaded!")
-
-    if prompt := st.chat_input("Your question"):
+    """Main function to run the chatbot app."""
+    st.set_page_config(page_title="Gemini Chatbot")
+    api_key = create_sidebar()  # Get the API key from the sidebar
+    st.title("💬 Gemini Chatbot")
+
+    if not api_key:
+        st.warning("👆 Please enter your Google API key in the sidebar to start.")
+        return
+
+    # Initialize Gemini model
+    genai.configure(api_key=api_key)
+    model = genai.GenerativeModel("gemini-2.0-flash")  # Use gemini-pro for chat
+
+    # Display chat history
+    st.markdown("### 💭 Chat")
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+
+    # Get user input and generate response
+    if prompt := st.chat_input("Say something..."):
         st.session_state.messages.append({"role": "user", "content": prompt})
-        logging.info(f"User input: {prompt}")
-
-        # Display conversation
-        for message in st.session_state.messages:
-            with st.chat_message(message["role"]):
-                st.write(message["content"])
-
-        if st.session_state.messages[-1]["role"] == "user":
-            with st.chat_message("assistant"):
-                start_time = time.time()
-                logging.info("Generating response")
-
-                with st.spinner("Writing..."):
-                    try:
-                        response_message = stream_chat(st.session_state.generator, st.session_state.messages)
-                        duration = time.time() - start_time
-                        response_with_duration = f"{response_message}\n\n⏱ Duration: {duration:.2f} seconds"
-                        st.session_state.messages.append({"role": "assistant", "content": response_with_duration})
-                        st.write(response_with_duration)
-                        logging.info(f"Response generated in {duration:.2f}s")
-
-                    except Exception as e:
-                        st.session_state.messages.append({"role": "assistant", "content": str(e)})
-                        st.error("An error occurred while generating the response.")
-                        logging.error(f"Error: {str(e)}")
+        with st.chat_message("user"):
+            st.markdown(prompt)
+
+        with st.chat_message("assistant"):
+            with st.spinner("Thinking..."):
+                try:
+                    response = model.generate_content(prompt)
+                    if response.text:
+                        st.markdown(response.text)
+                        st.session_state.messages.append({"role": "assistant", "content": response.text})
+                    else:
+                        st.markdown("Sorry, I couldn't get a response.") # error message
+                        st.session_state.messages.append({"role": "assistant", "content": "Sorry, I couldn't get a response."})
+                except Exception as e:
+                    st.error(f"An error occurred: {e}") # show the error to the user
+                    st.session_state.messages.append({"role": "assistant", "content": f"An error occurred: {e}"})
 
 if __name__ == "__main__":
-    main()
+    main()