modifying generation pipeline

app.py

@@ -1,22 +1,43 @@
 import streamlit as st
+
+from googletrans import Translator
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 # from huggingface_hub import snapshot_download
 
 page = st.sidebar.selectbox("Model ", ["Finetuned on News data", "Pretrained GPT2"])
+translator = Translator()
 
 def load_model(model_name):
     with st.spinner('Waiting for the model to load.....'):
         # snapshot_download('flax-community/Sinhala-gpt2')
-        model = AutoModelForCausalLM.from_pretrained(model_name)
         tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name, pad_token_id=tokenizer.eos_token_id)
     st.success('Model loaded!!')
     return model, tokenizer
 
 seed = st.sidebar.text_input('Starting text', 'ආයුබෝවන්')
-seq_num = st.sidebar.number_input('Number of sentences to generate ', 1, 20, 5)
-max_len = st.sidebar.number_input('Length of the sentence ', 5, 300, 100)
+seq_num = st.sidebar.number_input('Number of sequences to generate ', 1, 20, 5)
+max_len = st.sidebar.number_input('Length of a sequence ', 5, 300, 100)
 gen_bt = st.sidebar.button('Generate')
 
+def generate(model, tokenizer, seed, seq_num, max_len):
+    sentences = []
+    input_ids = tokenizer.encode(seed, return_tensors='pt')
+    beam_outputs = model.generate(
+        input_ids, 
+        do_sample=True, 
+        max_length=max_len, 
+        top_k=50, 
+        top_p=0.95, 
+        temperature=0.7,
+        num_return_sequences=seq_num,
+        no_repeat_ngram_size=2,
+        early_stopping=True
+    )
+    for beam_out in beam_outputs:
+        sentences.append(tokenizer.decode(beam_out, skip_special_tokens=True))
+    return sentences
+    
 if page == 'Pretrained GPT2':
     st.title('Sinhala Text generation with GPT2')
     st.markdown('A simple demo using Sinhala-gpt2 model trained during hf-flax week')
@@ -27,9 +48,14 @@ if page == 'Pretrained GPT2':
     if gen_bt:
         try:
             with st.spinner('Generating...'):
-                generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
-                seqs = generator(seed, max_length=max_len, num_return_sequences=seq_num)
-            st.write(seqs)
+                # generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
+                # seqs = generator(seed, max_length=max_len, num_return_sequences=seq_num)
+                seqs = generate(model, tokenizer, seed, seq_num, max_len)
+            for i, seq in enumerate(seqs):
+                st.info(f'Generated sequence {i+1}:')
+                st.write(seq)
+                st.info(f'English translation (by Google Translation):')
+                st.write(translator.translate(seq, src='si', dest='en').text)
         except Exception as e:
             st.exception(f'Exception: {e}')
 else:
@@ -43,13 +69,14 @@ else:
     if gen_bt:
         try:
             with st.spinner('Generating...'):
-                generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
-                seqs = generator(seed, max_length=max_len, num_return_sequences=seq_num)
-            st.write(seqs)
+                # generator = pipeline('text-generation', model=model, tokenizer=tokenizer)
+                # seqs = generator(seed, max_length=max_len, num_return_sequences=seq_num)
+                seqs = generate(model, tokenizer, seed, seq_num, max_len)
+            for i, seq in enumerate(seqs):
+                st.info(f'Generated sequence {i+1}:')
+                st.write(seq)
+                st.info(f'English translation (by Google Translation):')
+                st.write(translator.translate(seq, src='si', dest='en').text)
         except Exception as e:
             st.exception(f'Exception: {e}')
-           
-
 st.markdown('____________')
-st.markdown('by Keshan with Flax Community')
-

requirements.txt

@@ -1,5 +1,4 @@
 transformers
 streamlit
-jax
 torch
-flax
+googletrans==3.1.0a