Update app.py

app.py

@@ -1,43 +1,87 @@
-import streamlit as st
-from transformers import AutoModelForCausalLM, AutoTokenizer
+# import streamlit as st
+# 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
+# # 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()
+# 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.")
+# # 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"] = []
+# # 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")
+# # 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})
+# 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)
+#     # 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']}")
+import streamlit as st
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
+
+# Set up the app title
+st.title("🤖 Simple Chatbot")
 
-    # Append model response to history
-    st.session_state["conversation_history"].append({"role": "assistant", "content": response})
+# Initialize the Hugging Face pipeline
+@st.cache_resource
+def load_chatbot():
+    # Choose a model from Hugging Face Hub
+    model_name = "microsoft/DialoGPT-medium"  # Try other models like "gpt2", "facebook/blenderbot-400M-distill"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+
+chatbot = load_chatbot()
 
-# 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']}")
+# Initialize session state for chat history
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
+# Display chat history
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Chat input
+if prompt := st.chat_input("Type your message..."):
+    # Add user message to history
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    
+    # Generate bot response
+    with st.spinner("Thinking..."):
+        bot_response = chatbot(
+            prompt,
+            max_length=1000,
+            num_return_sequences=1,
+            pad_token_id=chatbot.tokenizer.eos_token_id
+        )[0]['generated_text']
+    
+    # Add bot response to history
+    st.session_state.messages.append({"role": "assistant", "content": bot_response})
+    
+    # Rerun to display new messages
+    st.rerun()