import streamlit as st
import logging
import time
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

# Set the logger to display only CRITICAL messages
logging.basicConfig(level=logging.CRITICAL)


# Cache the model and tokenizer to avoid reloading it every time
@st.experimental_singleton
def load_model():
    model_name = "Abbeite/trail_wl"  # Replace with your actual model name
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    return model, tokenizer

model, tokenizer = load_model()

# Function to generate text with the model
def generate_text(prompt):
    formatted_prompt = f"[INST] {prompt} [/INST]"  # Format the prompt according to your specification
    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_length=300)
    result = pipe(formatted_prompt)
    return result[0]['generated_text']

st.title("Interact with Your Model")

# User input
user_input = st.text_area("Enter your prompt:", "")

if st.button("Submit"):
    if user_input:
        generated_text = generate_text(user_input)
        for i in range(1, len(generated_text) + 1):
            # Update the output placeholder with a substring of the generated text
            output_placeholder.markdown(generated_text[:i])
            # Sleep to simulate typing
            time.sleep(0.1)  # Adjust the sleep time to speed up or slow down the "typing"
    else:
        st.write("Please enter a prompt to get started.")