flax-community/TamilLanguageDemos

1""" Script for streamlit demo2    @author: AbinayaM023"""4 5# Install necessary libraries6from transformers import AutoTokenizer, AutoModelWithLMHead, pipeline7import streamlit as st8import json9 10# Read the config11with open("config.json") as f:12    config = json.loads(f.read())13 14# Set page layout15st.set_page_config(16        page_title="Tamil Language Models",17        page_icon="U+270D",18        layout="wide",19        initial_sidebar_state="expanded"20    )21 22# Load the model23@st.cache(allow_output_mutation=True)24def load_model(model_name):25    with st.spinner('Waiting for the model to load.....'):26        model = AutoModelWithLMHead.from_pretrained(model_name)27        tokenizer = AutoTokenizer.from_pretrained(model_name)28    return model, tokenizer29 30# Side bar31img = st.sidebar.image("images/tamil_logo.jpg", width=300)32 33# Choose the model based on selection34st.sidebar.title("கதை சொல்லி!")35page = st.sidebar.selectbox(label="Select model", 36                            options=config["models"],37                            help="Select the model to generate the text")38data = st.sidebar.selectbox(label="Select data",39                            options=config[page],40                            help="Select the data on which the model is trained")41if page == "Text Generation" and data == "Oscar + IndicNLP":42    st.sidebar.markdown(43        "[Model tracking on wandb](https://wandb.ai/wandb/hf-flax-gpt2-tamil/runs/watdq7ib/overview?workspace=user-abinayam)",44         unsafe_allow_html=True45    )46    st.sidebar.markdown(47        "[Model card](https://huggingface.co/abinayam/gpt-2-tamil)",48        unsafe_allow_html=True49    )50elif page == "Text Generation" and data == "Oscar":51    st.sidebar.markdown(52        "[Model tracking on wandb](https://wandb.ai/abinayam/hf-flax-gpt-2-tamil/runs/1ddv4131/overview?workspace=user-abinayam)",53         unsafe_allow_html=True54    )55    st.sidebar.markdown(56        "[Model card](https://huggingface.co/flax-community/gpt-2-tamil)",57        unsafe_allow_html=True58    )59 60# Main page61st.title("Tamil Language Demos")62st.markdown(63    "Built as part of the Flax/Jax Community week, this demo uses [GPT2 trained on Oscar dataset](https://huggingface.co/flax-community/gpt-2-tamil) "64    "and [GPT2 trained on Oscar & IndicNLP dataset] (https://huggingface.co/abinayam/gpt-2-tamil) "65    "to show language generation!"66)67 68# Set default options for examples69prompts = config["examples"] + ["Custom"]70 71if page == 'Text Generation' and data == 'Oscar':72    st.header('Tamil text generation with GPT2')73    st.markdown('A simple demo using gpt-2-tamil model trained on Oscar dataset!')74    model, tokenizer = load_model(config[data])75elif page == 'Text Generation' and data == "Oscar + Indic Corpus":76    st.header('Tamil text generation with GPT2')77    st.markdown('A simple demo using gpt-2-tamil model trained on Oscar + IndicNLP dataset')78    model, tokenizer = load_model(config[data])79else:80    st.title('Tamil News classification with Finetuned GPT2')81    st.markdown('In progress')82 83if page == "Text Generation":84    # Set default options85    prompt = st.selectbox('Examples', prompts, index=0)86    if prompt == "Custom":87        prompt_box = "",88        text = st.text_input(89        'Add your custom text in Tamil',90        "",91        max_chars=1000)92    else:93        prompt_box = prompt94        text = st.text_input(95        'Selected example in Tamil',96        prompt,97        max_chars=1000)98    max_len = st.slider('Select length of the sentence to generate', 25, 300, 100)99    gen_bt = st.button('Generate')100 101    # Generate text102    if gen_bt:103            try:104                with st.spinner('Generating...'):105                    generator = pipeline('text-generation', model=model, tokenizer=tokenizer)106                    seqs = generator(prompt_box, max_length=max_len)[0]['generated_text']107                st.write(seqs)108            except Exception as e:109                st.exception(f'Exception: {e}')110