import streamlit as st
from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("BeardedMonster/SabiYarn")
repo_name = "BeardedMonster/SabiYarn-125M"
model = AutoModelForCausalLM.from_pretrained(repo_name, trust_remote_code=True)

# Define generation configuration
generation_config = GenerationConfig(
    max_length=100,
    num_beams=5,
    do_sample=True,
    temperature=0.9,
    top_k=50,
    top_p=0.95,
    repetition_penalty=2.0,
    length_penalty=1.7,
    early_stopping=True
)

# Streamlit app
st.title("SabiYarn-125M")
st.write("Generate text in multiple Nigerian languages using the SabiYarn model.")

# Text input
user_input = st.text_area("Enter your text here:", "")

if st.button("Generate"):
    if user_input:
        input_ids = tokenizer(user_input, return_tensors="pt")["input_ids"]
        output = model.generate(input_ids, generation_config=generation_config, max_new_tokens=50)
        input_len = len(input_ids[0])
        generated_text = tokenizer.decode(output[0][input_len:])
        st.write("### Generated Text")
        st.write(generated_text)
    else:
        st.write("Please enter some text to generate.")