import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import streamlit as st
import airllm
import os

from dotenv import load_dotenv

load_dotenv()


api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")

# Retrieve the API token from the environment variables
api_token = os.getenv("HUGGINGFACEHUB_API_TOKEN")

# Load GEMMA 27B model and tokenizer using the API token
tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b-it", use_auth_token=api_token)
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-9b-it",
    device_map="auto",
    torch_dtype=torch.bfloat16,
    use_auth_token=api_token
)
# Load the base version of the model
#model = AutoModelForCausalLM.from_pretrained(model_name)

# Initialize AirLLM
air_llm = airllm.AirLLM(model, tokenizer)

# Streamlit app configuration
st.set_page_config(
    page_title="Chatbot with GEMMA 27B and AirLLM",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded",
)

# App title
st.title("Conversational Chatbot with GEMMA 27B and AirLLM")

# Sidebar configuration
st.sidebar.header("Chatbot Configuration")
theme = st.sidebar.selectbox("Choose a theme", ["Default", "Dark", "Light"])

# Set theme based on user selection
if theme == "Dark":
    st.markdown(
        """
        <style>
        .reportview-container {
            background: #2E2E2E;
            color: #FFFFFF;
        }
        .sidebar .sidebar-content {
            background: #333333;
        }
        </style>
        """,
        unsafe_allow_html=True
    )
elif theme == "Light":
    st.markdown(
        """
        <style>
        .reportview-container {
            background: #FFFFFF;
            color: #000000;
        }
        .sidebar .sidebar-content {
            background: #F5F5F5;
        }
        </style>
        """,
        unsafe_allow_html=True
    )

# Chat input and output
user_input = st.text_input("You: ", "")
if st.button("Send"):
    if user_input:
        # Generate response using AirLLM
        response = air_llm.generate_response(user_input)
        st.text_area("Bot:", value=response, height=200, max_chars=None)
    else:
        st.warning("Please enter a message.")

# Footer
st.sidebar.markdown(
    """
    ### About
    This is a conversational chatbot built using the base version of the GEMMA 27B model and AirLLM.
    """
)