Update app.py

app.py

@@ -3,6 +3,9 @@ import random
 import time
 from ctransformers import AutoModelForCausalLM
 from datetime import datetime
+import whisper
+from transformers import VitsModel, AutoTokenizer
+import torch
 
 params = {
         "max_new_tokens":512,
@@ -13,15 +16,38 @@ params = {
         "batch_size": 8}
 
 
+whisper_model = whisper.load_model("small")
 llm = AutoModelForCausalLM.from_pretrained("Aspik101/trurl-2-7b-pl-instruct_GGML", model_type="llama")
+tts_model = VitsModel.from_pretrained("facebook/mms-tts-pol")
+tokenizer = AutoTokenizer.from_pretrained("facebook/mms-tts-pol")
+
 
 with gr.Blocks() as demo:
     chatbot = gr.Chatbot()
-    msg = gr.Textbox()
+    audio_input = gr.Audio(source="microphone", type="filepath", show_label=False)
+    submit_audio = gr.Button("Submit Audio")
     clear = gr.Button("Clear")
-
-    def user(user_message, history):
-        return "", history + [[user_message, None]]
+    audio_output = gr.Audio('temp_file.wav', label="Generated Audio (wav)", type='filepath', autoplay=False)
+    
+    def translate(audio):
+        print("__Wysyłam nagranie do whisper!")
+        transcription = whisper_model.transcribe(audio, language="pl")
+        return transcription["text"]
+    
+    def read_text(text):
+        print("Tutaj jest tekst to przeczytania!", text[-1][-1])
+        inputs = tokenizer(text[-1][-1], return_tensors="pt")
+        with torch.no_grad():
+            output = tts_model(**inputs).waveform.squeeze().numpy()
+        sf.write('temp_file.wav', output, tts_model.config.sampling_rate)
+        return 'temp_file.wav'
+    
+    def user(audio_data, history):
+        if audio_data:
+            user_message = translate(audio_data)
+            print("USER!:")
+            print("", history + [[user_message, None]])
+            return history + [[user_message, None]]
 
     def parse_history(hist):
         history_ = ""
@@ -33,9 +59,7 @@ with gr.Blocks() as demo:
 
     def bot(history):
         print(f"When: {datetime.today().strftime('%Y-%m-%d %H:%M:%S')}")
-        print("history: ",history)
-        prompt =  f"Jesteś AI assystentem. Odpowiadaj po polsku. {parse_history(history)}. <assistant>:"
-        print("prompt: ",prompt)
+        prompt = f"Jesteś AI assystentem. Odpowiadaj krótko i po polsku. {parse_history(history)}. <assistant>:"
         stream = llm(prompt, **params)
         history[-1][1] = ""
         answer_save = ""
@@ -45,12 +69,8 @@ with gr.Blocks() as demo:
             time.sleep(0.005)
             yield history
 
-        print("answer_save: ",answer_save)
-    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot, chatbot, chatbot
-    )
+    submit_audio.click(user, [audio_input, chatbot], [chatbot], queue=False).then(bot, chatbot, chatbot).then(read_text, chatbot, audio_output)
     clear.click(lambda: None, None, chatbot, queue=False)
 
-
 demo.queue()
-demo.launch()
+demo.launch(share=True)