Spaces:
Running
Running
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.") | |