import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM
import chromadb
from chromadb.config import Settings
from chromadb.utils import embedding_functions
from sentence_transformers import SentenceTransformer
import os
from huggingface_hub import login
from transformers import AutoModel

# Retrieve the API token from the environment variable
os.environ["HF_API_TOKEN"]  = os.getenv("HF_TOKEN")


# Set your Hugging Face API token as an environment variable
os.environ["HF_API_TOKEN"] = "your_huggingface_api_token"



# Load the Llama model using Hugging Face Transformers
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-chat-hf")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-7b-chat-hf")

# Initialize ChromaDB
client = chromadb.Client(Settings(chroma_db_impl="duckdb+parquet", persist_directory="./chroma_db"))

# Create a collection for storing supply chain and green environment data
collection = client.get_or_create_collection(
    name="supply_chain_green_environment",
    embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name="all-mpnet-base-v2"
    ),
)

# Initialize the sentence transformer for generating embeddings
embedding_model = SentenceTransformer("all-mpnet-base-v2")

# Streamlit app title
st.title("Supply Chain & Green Environment Chatbot")

# User input for questions
user_question = st.text_input("Enter your question:")

# Chat history
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# Process user input and generate response
if user_question:
    # Generate embedding for the user question
    question_embedding = embedding_model.encode(user_question).tolist()

    # Search for relevant information in the ChromaDB collection
    results = collection.query(
        query_embeddings=question_embedding,
        n_results=3,
    )

    # Construct the context for the Llama model
    context = ""
    for doc in results["documents"][0]:
        context += doc + "\n"

    # Generate response from the Llama model
    inputs = tokenizer(f"Context: {context}\n\nQuestion: {user_question}\n\nAnswer:", return_tensors="pt")
    outputs = model.generate(**inputs, max_length=256)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # Add user question and bot response to chat history
    st.session_state.chat_history.append({"user": user_question, "bot": response})

# Display chat history
for message in st.session_state.chat_history:
    st.write(f"**User:** {message['user']}")
    st.write(f"**Bot:** {message['bot']}")