Update app.py

app.py

@@ -1,71 +1,83 @@
 import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+import os
 
-# Load the tokenizer and model
-model_name = "meta-llama/Llama-3.2-1B-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-model.to('cuda' if torch.cuda.is_available() else 'cpu')  # Move model to GPU if available
+# Retrieve the Hugging Face token from environment variables
+HF_TOKEN = os.environ.get("GColab", None)
 
-# Streamlit app setup
-st.set_page_config(
-    page_title="AI Chatbot",
-    page_icon="🤖",
-    layout="wide",
-    initial_sidebar_state="expanded",
-)
+# Ensure the token is provided
+if HF_TOKEN is None:
+    st.error("Hugging Face token is not set. Please set the HF_TOKEN environment variable.")
+else:
+    # Load the tokenizer and model
+    model_name = "meta-llama/Llama-3.2-1B-Instruct"
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=HF_TOKEN)
+        model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=HF_TOKEN)
+        model.to('cuda' if torch.cuda.is_available() else 'cpu')  # Move model to GPU if available
 
-# Sidebar
-st.sidebar.title("Chatbot Settings")
-max_length = st.sidebar.slider("Response Length", min_value=50, max_value=512, value=100)
-temperature = st.sidebar.slider("Temperature", min_value=0.5, max_value=1.5, value=0.95)
-st.sidebar.markdown("Adjust the response settings for the chatbot.")
+        # Streamlit app setup
+        st.set_page_config(
+            page_title="AI Chatbot",
+            page_icon="🤖",
+            layout="wide",
+            initial_sidebar_state="expanded",
+        )
 
-# Main page
-st.title("🤖 Chat with AI")
-st.markdown("Welcome to the interactive AI chatbot. Start a conversation by typing below:")
+        # Sidebar
+        st.sidebar.title("Chatbot Settings")
+        max_length = st.sidebar.slider("Response Length", min_value=50, max_value=512, value=100)
+        temperature = st.sidebar.slider("Temperature", min_value=0.5, max_value=1.5, value=0.95)
+        st.sidebar.markdown("Adjust the response settings for the chatbot.")
 
-# Chat container
-if "conversation" not in st.session_state:
-    st.session_state.conversation = []
+        # Main page
+        st.title("🤖 Chat with AI")
+        st.markdown("Welcome to the interactive AI chatbot. Start a conversation by typing below:")
 
-def generate_response(user_input, history=[], max_new_tokens=512, temperature=0.95):
-    # Prepare the conversation history
-    conversation = [{"role": "user", "content": msg} for msg in history]
-    conversation.append({"role": "user", "content": user_input})
+        # Chat container
+        if "conversation" not in st.session_state:
+            st.session_state.conversation = []
 
-    # Tokenize the input
-    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
+        def generate_response(user_input, history=[], max_new_tokens=512, temperature=0.95):
+            # Prepare the conversation history
+            conversation = [{"role": "user", "content": msg} for msg in history]
+            conversation.append({"role": "user", "content": user_input})
 
-    # Generate response
-    outputs = model.generate(
-        input_ids,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        temperature=temperature,
-        eos_token_id=tokenizer.eos_token_id
-    )
+            # Tokenize the input
+            input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
 
-    # Decode and return the response
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response.split("assistant")[1]
+            # Generate response
+            outputs = model.generate(
+                input_ids,
+                max_new_tokens=max_new_tokens,
+                do_sample=True,
+                temperature=temperature,
+                eos_token_id=tokenizer.eos_token_id
+            )
 
-def add_to_conversation(user_input, bot_response):
-    st.session_state.conversation.append({"user": user_input, "bot": bot_response})
+            # Decode and return the response
+            response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            return response.split("assistant")[1]
 
-# User input
-user_input = st.text_input("You:", "")
-if st.button("Send") and user_input:
-    bot_response = generate_response(user_input, max_new_tokens=max_length, temperature=temperature)
-    add_to_conversation(user_input, bot_response)
-    user_input = ""  # Clear the input field
+        def add_to_conversation(user_input, bot_response):
+            st.session_state.conversation.append({"user": user_input, "bot": bot_response})
 
-# Display conversation
-for chat in st.session_state.conversation:
-    st.markdown(f"**You:** {chat['user']}")
-    st.markdown(f"**Bot:** {chat['bot']}")
+        # User input
+        user_input = st.text_input("You:", "")
+        if st.button("Send") and user_input:
+            bot_response = generate_response(user_input, max_new_tokens=max_length, temperature=temperature)
+            add_to_conversation(user_input, bot_response)
+            user_input = ""  # Clear the input field
 
-# Footer
-st.markdown("---")
-st.markdown("Developed by Akhil Sudhakaran")
+        # Display conversation
+        for chat in st.session_state.conversation:
+            st.markdown(f"**You:** {chat['user']}")
+            st.markdown(f"**Bot:** {chat['bot']}")
+
+        # Footer
+        st.markdown("---")
+        st.markdown("Developed by Akhil Sudhakaran")
+
+    except Exception as e:
+        st.error(f"An error occurred: {e}")