Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from transformers import ( | |
| AutoModelForCausalLM, | |
| AutoModelForSeq2SeqLM, | |
| AutoTokenizer, | |
| GenerationConfig, | |
| ) | |
| def load_tokenizer_model(name_or_path: str, model_type: str, model_auth_token: str): | |
| model_auth_token = None if model_auth_token == "" else model_auth_token | |
| tokenizer = AutoTokenizer.from_pretrained(*name_or_path.split(","), use_auth_token=model_auth_token) | |
| if model_type == "seq2seq": | |
| model = AutoModelForSeq2SeqLM.from_pretrained(*name_or_path.split(","), use_auth_token=model_auth_token) | |
| elif model_type == "causal": | |
| model = AutoModelForCausalLM.from_pretrained(*name_or_path.split(","), use_auth_token=model_auth_token) | |
| else: | |
| raise ValueError("model_type must be one of 'seq2seq' or 'causal'") | |
| return tokenizer, model | |
| def main(): | |
| st.title("Huggingface Transformers Demo") | |
| with st.form("model_form"): | |
| model_type = st.selectbox("Select Model Type", ["seq2seq", "causal"]) | |
| model_name_or_path = st.text_input("Model Name or Path") | |
| model_auth_token = st.text_input("Model Auth Token") | |
| input_text = st.text_area("Input Text") | |
| col1, col2, col3 = st.columns(3) | |
| user_gen_config = {} | |
| with col1: | |
| user_gen_config["min_length"] = st.number_input("Min Length", value=10, min_value=1, max_value=1000, step=1) | |
| user_gen_config["max_length"] = st.number_input("Max Length", value=50, min_value=1, max_value=1000, step=1) | |
| user_gen_config["top_k"] = st.number_input("Top K", value=50, min_value=1, max_value=100, step=1) | |
| with col2: | |
| user_gen_config["num_beams"] = st.number_input("Num Beams", value=1, min_value=1, max_value=100, step=1) | |
| user_gen_config["top_p"] = st.number_input("Top P", value=1.0, min_value=0.0, max_value=100.0, step=0.1) | |
| user_gen_config["repetition_penalty"] = st.number_input("Repetition Penalty", value=1.0, min_value=0.0, | |
| max_value=100.0, step=0.1) | |
| with col3: | |
| user_gen_config["temperature"] = st.number_input("Temperature", value=1.0, min_value=0.0, max_value=100.0, step=0.1) | |
| user_gen_config["do_sample"] = st.checkbox("Do Sample", value=False) | |
| user_gen_config["early_stopping"] = st.checkbox("Early Stopping", value=True) | |
| submitted = st.form_submit_button("Submit") | |
| if submitted: | |
| tokenizer, model = load_tokenizer_model(model_name_or_path, model_type, model_auth_token) | |
| gen_config = GenerationConfig.from_model_config(model.config) | |
| for k,v in user_gen_config.items(): | |
| setattr(gen_config, k, v) | |
| input_ids = tokenizer.encode(input_text, return_tensors="pt") | |
| output_ids = model.generate(input_ids, generation_config=gen_config) | |
| output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True) | |
| st.write(output_text) | |
| if __name__ == "__main__": | |
| main() | |