import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer

@st.cache(allow_output_mutation=True)
def load_model():
    model_id = "Tech-Meld/Hajax_Chat_1.0"
    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(model_id)
    return model, tokenizer

def get_response(input_text, model, tokenizer):
    inputs = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt')
    outputs = model.generate(inputs, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    response = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
    return response

def main():
    model, tokenizer = load_model()
    st.title("Chat with AI")
    input_text = st.text_input("You: ", "")
    if st.button("Send"):
        response = get_response(input_text, model, tokenizer)
        st.text_area("AI: ", response)

if __name__ == "__main__":
    main()