Update app.py

app.py

@@ -1,43 +1,15 @@
-
-import streamlit as st
-import os
-hf_token = os.environ.get("HF_TOKEN")
-if hf_token:
-    # Use the token
-    from huggingface_hub import login
-    login(token = hf_token)
-    #your code that requires the token.
-else:
-    print("HF_TOKEN environment variable not set.")
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
-import torch
-if torch.cuda.is_available():
-    device = torch.device("cuda")
-    print(f"CUDA is available. Using GPU: {torch.cuda.get_device_name(0)}") #prints GPU name
-    print(f"Number of GPUs available: {torch.cuda.device_count()}") #prints number of gpus.
-    print(f"Current GPU device: {torch.cuda.current_device()}")#prints current gpu id.
-else:
-    device = torch.device("cpu")
-    print("CUDA is not available. Using CPU.")
-
-print(f"Using device: {device}")
-
+# Initialize model and tokenizer (load only once)
 @st.cache_resource
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
-    model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it")
+    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+    model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
     return tokenizer, model
 
 tokenizer, model = load_model()
 
 # Function to generate chatbot response
-def generate_response(prompt, chat_history_ids=None):
-    inputs = tokenizer.encode(prompt + tokenizer.eos_token, return_tensors="pt")
-
-    if chat_history_ids is None:
-        chat_history_ids = None
-    else:
-        chat_history_ids = torch.tensor(chat_history_ids)
+def generate_response(prompt, chat_history=""):
+    inputs = tokenizer.encode(chat_history + prompt + tokenizer.eos_token, return_tensors="pt")
 
     # generate a response while limiting the total chat history to 1000 tokens,
     chat_history_ids = model.generate(
@@ -47,11 +19,10 @@ def generate_response(prompt, chat_history_ids=None):
         temperature=0.7,
         top_k=50,
         top_p=0.95,
-        chat_history_ids = chat_history_ids
     )
 
     response = tokenizer.decode(chat_history_ids[:, inputs.shape[-1]:][0], skip_special_tokens=True)
-    return response, chat_history_ids.tolist()
+    return response
 
 # Streamlit app
 st.title("Simple Chatbot")
@@ -59,8 +30,8 @@ st.title("Simple Chatbot")
 # Initialize chat history
 if "messages" not in st.session_state:
     st.session_state.messages = []
-if "chat_history_ids" not in st.session_state:
-    st.session_state.chat_history_ids = None
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = ""
 
 # Display chat messages from history on app rerun
 for message in st.session_state.messages:
@@ -79,7 +50,7 @@ if prompt := st.chat_input("What is up?"):
     with st.chat_message("assistant"):
         message_placeholder = st.empty()
         full_response = ""
-        response, st.session_state.chat_history_ids = generate_response(prompt, st.session_state.chat_history_ids)
+        response = generate_response(prompt, st.session_state.chat_history)
 
         # Simulate stream of responses with milliseconds delay
         import time
@@ -92,3 +63,5 @@ if prompt := st.chat_input("What is up?"):
 
     # Add assistant response to chat history
     st.session_state.messages.append({"role": "assistant", "content": full_response})
+    #update the chat history.
+    st.session_state.chat_history += prompt + tokenizer.eos_token + response + tokenizer.eos_token