Update app.py

app.py

@@ -1,99 +1,42 @@
 import streamlit as st
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-
-# Configure the Hugging Face API key (no need to pass it in the pipeline call)
-HF_API_KEY = st.secrets['huggingface_api_key']
-
-# Ensure you're logged in using the Hugging Face CLI if using private models
-# huggingface-cli login
-
-# Initialize the Hugging Face model and tokenizer
-model_name = 'gpt2-medium'  # or another GPT-2 version you want to use
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Initialize the text generation pipeline
-generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
-
-# Function to get response from the Hugging Face model
-def get_chatbot_response(user_input):
-    try:
-        # Generate the response using the Hugging Face model
-        response = generator(user_input, max_length=100, num_return_sequences=1)
-        return response[0]['generated_text']
-    except Exception as e:
-        return f"Error: {str(e)}"
-
-# Streamlit interface
-st.set_page_config(page_title="Smart ChatBot", layout="centered")
-
-# Custom CSS for chat bubbles with full width and emojis
-st.markdown("""
-    <style>
-    .chat-container {
-        display: flex;
-        flex-direction: column;
-        width: 100%;
-    }
-    .chat-bubble {
-        width: 100%;
-        padding: 15px;
-        margin: 10px 0;
-        border-radius: 10px;
-        font-size: 18px;
-        color: white;
-        display: inline-block;
-        line-height: 1.5;
-    }
-    .user-bubble {
-        background: #6a82fb; /* Soft blue */
-        align-self: flex-end;
-        border-radius: 10px 10px 10px 10px;
-    }
-    .bot-bubble {
-        background: #fc5c7d; /* Soft pink */
-        align-self: flex-start;
-        border-radius: 10px 10px 10px 10px;
-    }
-    .chat-header {
-        font-size: 35px;
-        font-weight: bold;
-        margin-bottom: 20px;
-        color: #3d3d3d;
-    }
-    .emoji {
-        font-size: 22px;
-        margin-right: 10px;
-    }
-    </style>
-""", unsafe_allow_html=True)
-
-# Chat header and intro
-st.markdown('<div class="chat-header">AI Chatbot - Your Companion 💻</div>', unsafe_allow_html=True)
-st.write("Powered by Hugging Face AI for smart, engaging conversations. 🤖")
-
-# Initialize session state for conversation history if not already initialized
-if "history" not in st.session_state:
-    st.session_state["history"] = []
-
-# Create the chat form
-with st.form(key="chat_form", clear_on_submit=True):
-    user_input = st.text_input("Your message here... ✍️", max_chars=2000, label_visibility="collapsed")
-    submit_button = st.form_submit_button("Send 🚀")
-
-    if submit_button:
-        if user_input:
-            # Get response from the chatbot
-            response = get_chatbot_response(user_input)
-            # Store user input and bot response in session state history
-            st.session_state.history.append((user_input, response))
-        else:
-            st.warning("Please Enter A Prompt 😅")
-
-# Display chat history
-if st.session_state["history"]:
-    st.markdown('<div class="chat-container">', unsafe_allow_html=True)
-    for user_input, response in st.session_state["history"]:
-        st.markdown(f'<div class="chat-bubble user-bubble"><span class="emoji">👤</span>You: {user_input}</div>', unsafe_allow_html=True)
-        st.markdown(f'<div class="chat-bubble bot-bubble"><span class="emoji">🤖</span>Bot: {response}</div>', unsafe_allow_html=True)
-    st.markdown('</div>', unsafe_allow_html=True)
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+# Load the model and tokenizer
+@st.cache_resource
+def load_model_and_tokenizer():
+    model_name = "microsoft/DialoGPT-medium"  # Replace with your chosen model
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return tokenizer, model
+
+tokenizer, model = load_model_and_tokenizer()
+
+# Streamlit App
+st.title("General Chatbot")
+st.write("A chatbot powered by an open-source model from Hugging Face.")
+
+# Initialize the conversation
+if "conversation_history" not in st.session_state:
+    st.session_state["conversation_history"] = []
+
+# Input box for user query
+user_input = st.text_input("You:", placeholder="Ask me anything...", key="user_input")
+
+if st.button("Send") and user_input:
+    # Append user input to history
+    st.session_state["conversation_history"].append({"role": "user", "content": user_input})
+    
+    # Tokenize and generate response
+    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
+    chat_history_ids = model.generate(input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
+    response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
+
+    # Append model response to history
+    st.session_state["conversation_history"].append({"role": "assistant", "content": response})
+
+# Display the conversation
+for message in st.session_state["conversation_history"]:
+    if message["role"] == "user":
+        st.write(f"**You:** {message['content']}")
+    else:
+        st.write(f"**Bot:** {message['content']}")