Update app.py

app.py

@@ -1,31 +1,49 @@
+import streamlit as st
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import os
+from huggingface_hub import login
+
+# Login with HF_TOKEN (if available)
+hf_token = os.environ.get("HF_TOKEN")
+if hf_token:
+    try:
+        login(token=hf_token, add_to_git_credential=False)
+        st.success("Hugging Face login successful!")
+    except Exception as e:
+        st.error(f"Hugging Face login failed: {e}")
+else:
+    st.warning("HF_TOKEN environment variable not set. Some features may be limited.")
+
 # Initialize model and tokenizer (load only once)
 @st.cache_resource
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
-    model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
+    model_name = "google/gemma-2b-it"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
     return tokenizer, model
 
 tokenizer, model = load_model()
 
 # Function to generate chatbot response
 def generate_response(prompt, chat_history=""):
-    inputs = tokenizer.encode(chat_history + prompt + tokenizer.eos_token, return_tensors="pt")
+    inputs = tokenizer.encode(chat_history + prompt, return_tensors="pt")
 
-    # generate a response while limiting the total chat history to 1000 tokens,
-    chat_history_ids = model.generate(
-        inputs, max_length=1000,
+    # Generate a response
+    outputs = model.generate(
+        inputs,
+        max_length=1000,
         pad_token_id=tokenizer.eos_token_id,
-        no_repeat_ngram_size=3,
         temperature=0.7,
         top_k=50,
         top_p=0.95,
     )
 
-    response = tokenizer.decode(chat_history_ids[:, inputs.shape[-1]:][0], skip_special_tokens=True)
+    response = tokenizer.decode(outputs[0][inputs.shape[-1]:], skip_special_tokens=True)
     return response
 
 # Streamlit app
-st.title("Simple Chatbot")
+st.title("Gemma-2b-it Chatbot")
 
 # Initialize chat history
 if "messages" not in st.session_state:
@@ -64,4 +82,4 @@ 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
+    st.session_state.chat_history += prompt + response