import streamlit as st
from transformers import BertTokenizer, TFBertForSequenceClassification
import tensorflow as tf
import numpy as np
import requests
import os

# Ensure PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION is set to "python"
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"

# Paths to your models hosted on Hugging Face
basic_model_url = "https://huggingface.co/anshupatel4298/bert-chatbot-model/resolve/main/basic_chatbot_model.h5"
local_model_path = "basic_chatbot_model.h5"
bert_model_name = "anshupatel4298/bert-chatbot-model/bert_model"

# Define the model architecture
def create_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Dense(128, activation='relu', input_shape=(100,)),  # Adjust input shape accordingly
        tf.keras.layers.Dropout(0.5),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dropout(0.5),
        tf.keras.layers.Dense(10, activation='softmax')  # Adjust the number of classes accordingly
    ])
    return model

# Create the model
basic_model = create_model()

# Download the Basic Model weights from the URL if not already downloaded
if not os.path.exists(local_model_path):
    response = requests.get(basic_model_url)
    with open(local_model_path, 'wb') as f:
        f.write(response.content)

# Load the weights into the model
basic_model.load_weights(local_model_path)

# Load BERT Model and Tokenizer
bert_model = TFBertForSequenceClassification.from_pretrained(bert_model_name)
bert_tokenizer = BertTokenizer.from_pretrained(bert_model_name)

# Set your MAX_SEQUENCE_LENGTH here
MAX_SEQUENCE_LENGTH = 100

# Streamlit UI
st.sidebar.title("Select Model")
model_choice = st.sidebar.selectbox("Choose a model:", ["Basic Model", "BERT Model"])

st.title("Chatbot Interface")

user_input = st.text_input("You:")
if st.button("Send"):
    if user_input:
        if model_choice == "Basic Model":
            # Preprocess input for basic model
            tokenized_input = tf.keras.preprocessing.text.Tokenizer().texts_to_sequences([user_input])
            input_data = tf.keras.preprocessing.sequence.pad_sequences(tokenized_input, maxlen=MAX_SEQUENCE_LENGTH)
            prediction = basic_model.predict(input_data)
            response = np.argmax(prediction, axis=-1)[0]
        else:
            # Preprocess input for BERT model
            inputs = bert_tokenizer(user_input, return_tensors="tf", max_length=MAX_SEQUENCE_LENGTH, truncation=True, padding="max_length")
            outputs = bert_model(**inputs)
            response = tf.argmax(outputs.logits, axis=-1).numpy()[0]

        st.write(f"Bot: {response}")