Update app.py

app.py

@@ -5,11 +5,22 @@ from transformers import HubertForCTC, Wav2Vec2Processor , pipeline , Wav2Vec2Fo
 import torch
 import spacy 
 from spacy import displacy
+import en_core_web_sm
+import spacy.cli
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+import nltk
+from nltk import tokenize
+nltk.download('punkt')
+import spacy_streamlit
+
 
 st.title('Audio-to-Text')
 
 audio_file = st.file_uploader('Upload Audio' , type=['wav' , 'mp3','m4a'])
 
+st.title( 'Please select any of the NLP tasks')
+
+
 if st.button('Trascribe Audio'):
     if audio_file is not None:
         processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
@@ -19,7 +30,9 @@ if st.button('Trascribe Audio'):
         logits = model(input_values).logits
         predicted_ids = torch.argmax(logits, dim=-1)
         text = processor.batch_decode(predicted_ids)
-        st.write(text) 
+        summary_list = [str(sentence) for sentence in text]
+        result = ' '.join(summary_list)
+        st.markdown(result) 
     else:
         st.error('please upload the audio file')
 
@@ -33,8 +46,10 @@ if st.button('Summarize'):
         logits = model(input_values).logits
         predicted_ids = torch.argmax(logits, dim=-1)
         text = processor.batch_decode(predicted_ids)
+        summary_list = [str(sentence) for sentence in text]
+        result = ' '.join(summary_list)
         summarize = pipeline("summarization")
-        st.write(summarize(text))
+        st.markdown(summarize(result)[0]['summary_text'])
 
 if st.button('sentiment-analysis'):
         processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
@@ -44,8 +59,10 @@ if st.button('sentiment-analysis'):
         logits = model(input_values).logits
         predicted_ids = torch.argmax(logits, dim=-1)
         text = processor.batch_decode(predicted_ids)
+        summary_list = [str(sentence) for sentence in text]
+        result = ' '.join(summary_list)
         nlp_sa = pipeline("sentiment-analysis")
-        st.write(nlp_sa(text))
+        st.markdown(nlp_sa(result))
 
 if st.button('Name'):
         processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
@@ -55,7 +72,41 @@ if st.button('Name'):
         logits = model(input_values).logits
         predicted_ids = torch.argmax(logits, dim=-1)
         text = processor.batch_decode(predicted_ids)
-        str  = ''.join(text)
-        trf = spacy.load('en_core_web_sm')
-        doc=trf(str)
-        print(displacy.render(doc,style='ent'))
+        summary_list = [str(sentence) for sentence in text]
+        result = ' '.join(summary_list)
+        nlp = spacy.load('en_core_web_sm')
+        doc=nlp(result)
+        spacy_streamlit.visualize_ner(doc, labels=nlp.get_pipe("ner").labels, title= "List of Entities")
+
+
+tokenizer = AutoTokenizer.from_pretrained("t5-base")
+
+@st.cache(allow_output_mutation=True)
+def load_model():
+    model = AutoModelForSeq2SeqLM.from_pretrained("t5-base")
+    return model
+
+model1 = load_model()
+
+st.subheader('Select your source and target language below.')
+source_lang = st.selectbox("Source language",['English'])
+target_lang = st.selectbox("Target language",['German','French'])
+
+
+if st.button('Translate'):
+        processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
+        model = HubertForCTC.from_pretrained("facebook/hubert-large-ls960-ft")  
+        speech, rate = librosa.load(audio_file, sr=16000)
+        input_values = processor(speech, return_tensors="pt", padding="longest", sampling_rate=rate).input_values
+        logits = model(input_values).logits
+        predicted_ids = torch.argmax(logits, dim=-1)
+        text = processor.batch_decode(predicted_ids)
+        summary_list = [str(sentence) for sentence in text]
+        result = ' '.join(summary_list)
+        prefix = 'translate '+str(source_lang)+' to '+str(target_lang)
+        sentence_token =  tokenize.sent_tokenize(result)
+        output = tokenizer([prefix+sentence for sentence in sentence_token], padding=True, return_tensors="pt")
+        translated_id = model1.generate(output["input_ids"], attention_mask=output['attention_mask'], max_length=100)
+        translated_word = tokenizer.batch_decode(translated_id, skip_special_tokens=True)
+        st.subheader('Translated Text')
+        st.write(' '.join(translated_word))