agnixcode/multimodel_chatbot

1import os2import streamlit as st3from groq import Groq4from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq, pipeline5from espnet2.bin.tts_inference import Text2Speech6import soundfile as sf7from pydub import AudioSegment8import io9from streamlit_webrtc import webrtc_streamer, WebRtcMode, AudioProcessorBase10import av11import numpy as np12import nltk13 14# Download NLTK data15nltk.download("averaged_perceptron_tagger")16nltk.download("cmudict")17 18# Load Groq API key from environment secrets19GROQ_API_KEY = os.getenv("GROQ_API_KEY")20if not GROQ_API_KEY:21    st.error("Groq API key not found. Please add it as a secret.")22    st.stop()23 24# Initialize Groq client25groq_client = Groq(api_key=GROQ_API_KEY)26 27# Load models28@st.cache_resource  # Use st.cache_resource for caching models29def load_models():30    # Speech-to-Text31    processor = AutoProcessor.from_pretrained("openai/whisper-small")32    stt_model = AutoModelForSpeechSeq2Seq.from_pretrained("openai/whisper-small")33    stt_pipe = pipeline(34        "automatic-speech-recognition",35        model=stt_model,36        tokenizer=processor.tokenizer,37        feature_extractor=processor.feature_extractor,38        return_timestamps=True  # Enable timestamps for long-form audio39    )40 41    # Text-to-Speech42    tts_model = Text2Speech.from_pretrained("espnet/espnet_tts_vctk_espnet_spk_voxceleb12_rawnet")43 44    return stt_pipe, tts_model45 46stt_pipe, tts_model = load_models()47 48# Audio recorder49class AudioRecorder(AudioProcessorBase):50    def __init__(self):51        self.audio_frames = []52 53    def recv(self, frame: av.AudioFrame) -> av.AudioFrame:54        self.audio_frames.append(frame.to_ndarray())55        return frame56 57# Streamlit app58st.title("Voice and Text Chatbot")59 60# Sidebar for mode selection61mode = st.sidebar.radio("Select Mode", ["Text Chatbot", "Voice Chatbot"])62 63if mode == "Text Chatbot":64    # Text Chatbot65    st.header("Text Chatbot")66    user_input = st.text_input("Enter your message:")67 68    if user_input:69        try:70            # Generate response using Groq API71            chat_completion = groq_client.chat.completions.create(72                messages=[{"role": "user", "content": user_input}],73                model="mixtral-8x7b-32768",74                temperature=0.5,75                max_tokens=102476            )77            response = chat_completion.choices[0].message.content78            st.write("Generated Response:", response)79 80            # Convert response to speech81            speech, *_ = tts_model(response, spembs=tts_model.spembs[0])  # Use the first speaker embedding82            sf.write("response.wav", speech, 22050)83            st.audio("response.wav")84        except Exception as e:85            st.error(f"Error generating response: {e}")86 87elif mode == "Voice Chatbot":88    # Voice Chatbot89    st.header("Voice Chatbot")90 91    # Audio recorder92    st.write("Record your voice:")93    webrtc_ctx = webrtc_streamer(94        key="audio-recorder",95        mode=WebRtcMode.SENDONLY,96        audio_processor_factory=AudioRecorder,97        media_stream_constraints={"audio": True, "video": False},98    )99 100    if webrtc_ctx.audio_processor:101        st.write("Recording... Press 'Stop' to finish recording.")102 103        # Save recorded audio to a WAV file104        if st.button("Stop and Process Recording"):105            audio_frames = webrtc_ctx.audio_processor.audio_frames106            if audio_frames:107                # Combine audio frames into a single array108                audio_data = np.concatenate(audio_frames)109                # Save as WAV file110                sf.write("recorded_audio.wav", audio_data, samplerate=16000)111                st.success("Recording saved as recorded_audio.wav")112 113                # Process the recorded audio114                speech, _ = sf.read("recorded_audio.wav")115                output = stt_pipe(speech)  # Transcribe with timestamps116 117                # Display the full transcribed text118                st.write("Transcribed Text:", output['text'])119 120                # Display the text with timestamps (optional)121                if 'chunks' in output:122                    st.write("Transcribed Text with Timestamps:")123                    for chunk in output['chunks']:124                        st.write(f"{chunk['timestamp'][0]:.2f} - {chunk['timestamp'][1]:.2f}: {chunk['text']}")125 126                # Generate response using Groq API127                try:128                    chat_completion = groq_client.chat.completions.create(129                        messages=[{"role": "user", "content": output['text']}],130                        model="mixtral-8x7b-32768",131                        temperature=0.5,132                        max_tokens=1024133                    )134                    response = chat_completion.choices[0].message.content135                    st.write("Generated Response:", response)136 137                    # Convert response to speech138                    speech, *_ = tts_model(response, spembs=tts_model.spembs[0])  # Use the first speaker embedding139                    sf.write("response.wav", speech, 22050)140                    st.audio("response.wav")141                except Exception as e:142                    st.error(f"Error generating response: {e}")143            else:144                st.error("No audio recorded. Please try again.")