Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import requests | |
| import uuid | |
| import os | |
| from typing import Optional | |
| import tempfile | |
| from pydub import AudioSegment | |
| import re | |
| ASR_API = "http://astarwiz.com:9998/asr" | |
| TTS_SPEAK_SERVICE = 'http://astarwiz.com:9603/speak' | |
| TTS_WAVE_SERVICE = 'http://astarwiz.com:9603/wave' | |
| LANGUAGE_MAP = { | |
| "en": "English", | |
| "ma": "Malay", | |
| "ta": "Tamil", | |
| "zh": "Chinese" | |
| } | |
| # Add a password for developer mode | |
| DEVELOPER_PASSWORD = os.getenv("DEV_PWD") | |
| # Add this constant for the RapidAPI key | |
| #RAPID_API_KEY = os.getenv("RAPID_API_KEY") | |
| RAPID_API_KEY = os.getenv("RAPID_API_KEY") | |
| # Add this constant for available speakers | |
| AVAILABLE_SPEAKERS = { | |
| "en": ["MS"], | |
| "ma": ["msFemale"], | |
| "ta": ["ta_female1"], | |
| "zh": ["childChinese2"] | |
| } | |
| def fetch_youtube_id(youtube_url: str) -> str: | |
| if 'v=' in youtube_url: | |
| return youtube_url.split("v=")[1].split("&")[0] | |
| elif 'youtu.be/' in youtube_url: | |
| return youtube_url.split("youtu.be/")[1] | |
| elif 'shorts' in youtube_url: | |
| return youtube_url.split("/")[-1] | |
| else: | |
| raise Exception("Unsupported URL format") | |
| def download_youtube_audio(youtube_url: str, output_dir: Optional[str] = None) -> Optional[str]: | |
| video_id = fetch_youtube_id(youtube_url) | |
| if not video_id: | |
| return None | |
| if output_dir is None: | |
| output_dir = tempfile.gettempdir() | |
| output_filename = os.path.join(output_dir, f"{video_id}.mp3") | |
| if os.path.exists(output_filename): | |
| return output_filename # Return if the file already exists | |
| url = "https://youtube86.p.rapidapi.com/api/youtube/links" | |
| headers = { | |
| 'Content-Type': 'application/json', | |
| 'x-rapidapi-host': 'youtube86.p.rapidapi.com', | |
| 'x-rapidapi-key': RAPID_API_KEY | |
| } | |
| data = { | |
| "url": youtube_url | |
| } | |
| response = requests.post(url, headers=headers, json=data) | |
| print('Fetched audio links') | |
| if response.status_code == 200: | |
| result = response.json() | |
| for url in result[0]['urls']: | |
| if url.get('isBundle'): | |
| audio_url = url['url'] | |
| extension = url['extension'] | |
| audio_response = requests.get(audio_url) | |
| if audio_response.status_code == 200: | |
| temp_filename = os.path.join(output_dir, f"{video_id}.{extension}") | |
| with open(temp_filename, 'wb') as audio_file: | |
| audio_file.write(audio_response.content) | |
| # Convert to MP3 and downsample to 16000 Hz | |
| audio = AudioSegment.from_file(temp_filename, format=extension) | |
| audio = audio.set_frame_rate(16000) | |
| audio.export(output_filename, format="mp3", parameters=["-ar", "16000"]) | |
| os.remove(temp_filename) # Remove the temporary file | |
| return output_filename # Return the final MP3 filename | |
| return None # Return None if no successful download occurs | |
| else: | |
| print("Error:", response.status_code, response.text) | |
| return None # Return None on failure | |
| punctuation_marks = r'([\.!?！？。])' | |
| """ | |
| def split_text_with_punctuation(text): | |
| # Split the text using the punctuation marks, keeping the punctuation marks | |
| split_text = re.split(punctuation_marks, text) | |
| # Combine each punctuation mark with the preceding segment | |
| combined_segments = [] | |
| for i in range(0, len(split_text) - 1, 2): | |
| combined_segments.append(split_text[i] + split_text[i + 1]) | |
| # If there's any remaining text after the last punctuation, append it as well | |
| if len(split_text) % 2 != 0 and split_text[-1]: | |
| combined_segments.append(split_text[-1]) | |
| return combined_segments | |
| """ | |
| def split_text_with_punctuation(text): | |
| # Split the text using the punctuation marks, keeping the punctuation marks | |
| split_text = re.split(punctuation_marks, text) | |
| # Combine each punctuation mark with the preceding segment | |
| combined_segments = [] | |
| # Loop through the split text in steps of 2 | |
| for i in range(0, len(split_text) - 1, 2): | |
| combined_segments.append(split_text[i] + split_text[i + 1]) | |
| # Handle any remaining text that doesn't have a punctuation following it | |
| if len(split_text) % 2 != 0 and split_text[-1]: | |
| combined_segments.append(split_text[-1]) | |
| # Split any segment that exceeds 50 words | |
| final_segments = [] | |
| for segment in combined_segments: | |
| words = segment.split() # Split each segment into words | |
| if len(words) > 50: | |
| # Split the segment into chunks of no more than 50 words | |
| for j in range(0, len(words), 50): | |
| final_segments.append(' '.join(words[j:j+50])) | |
| else: | |
| final_segments.append(segment) | |
| return [segment for segment in final_segments if segment] # Filter out empty strings | |
| def inference_via_llm_api(input_text, min_new_tokens=2, max_new_tokens=64): | |
| print(input_text) | |
| one_vllm_input = f"<|im_start|>system\nYou are a translation expert.<|im_end|>\n<|im_start|>user\n{input_text}<|im_end|>\n<|im_start|>assistant" | |
| vllm_api = 'http://astarwiz.com:2333/' + "v1/completions" | |
| data = { | |
| "prompt": one_vllm_input, | |
| 'model': "./Edu-4B-NewTok-V2-20240904/", | |
| 'min_tokens': min_new_tokens, | |
| 'max_tokens': max_new_tokens, | |
| 'temperature': 0.1, | |
| 'top_p': 0.75, | |
| 'repetition_penalty': 1.1, | |
| "stop_token_ids": [151645, ], | |
| } | |
| response = requests.post(vllm_api, headers={"Content-Type": "application/json"}, json=data).json() | |
| print(response) | |
| if "choices" in response.keys(): | |
| return response["choices"][0]['text'].strip() | |
| else: | |
| return "The system got some error during vLLM generation. Please try it again." | |
| def transcribe_and_speak(audio, source_lang, target_lang, youtube_url=None, target_speaker=None): | |
| if youtube_url: | |
| audio = download_youtube_audio(youtube_url) | |
| if not audio: | |
| return "Failed to download YouTube audio.", None, None | |
| if not audio: | |
| return "Please provide an audio input or a valid YouTube URL.", None, None | |
| # ASR | |
| file_id = str(uuid.uuid4()) | |
| files = {'file': open(audio, 'rb')} | |
| data = { | |
| 'language': 'ms' if source_lang == 'ma' else source_lang, | |
| 'model_name': 'whisper-large-v2-local-cs', | |
| 'with_timestamp': False | |
| } | |
| asr_response = requests.post(ASR_API, files=files, data=data) | |
| print(asr_response.json()) | |
| if asr_response.status_code == 200: | |
| transcription = asr_response.json()['text'] | |
| else: | |
| return "ASR failed", None, None | |
| split_result = split_text_with_punctuation(transcription) | |
| translate_segments=[] | |
| for segment in split_result: | |
| translation_prompt = f"Translate the following text from {LANGUAGE_MAP[source_lang]} to {LANGUAGE_MAP[target_lang]}: {segment}" | |
| translated_seg_txt = inference_via_llm_api(translation_prompt) | |
| translate_segments.append(translated_seg_txt) | |
| print(f"Translation: {translated_seg_txt}") | |
| translated_text = " ".join(translate_segments) | |
| # TTS | |
| tts_params = { | |
| 'language': target_lang, | |
| 'speed': 1.1, | |
| 'speaker': target_speaker or AVAILABLE_SPEAKERS[target_lang][0], # Use the first speaker as default | |
| 'text': translated_text | |
| } | |
| tts_response = requests.get(TTS_SPEAK_SERVICE, params=tts_params) | |
| if tts_response.status_code == 200: | |
| audio_file = tts_response.text.strip() | |
| audio_url = f"{TTS_WAVE_SERVICE}?file={audio_file}" | |
| return transcription, translated_text, audio_url | |
| else: | |
| return transcription, translated_text, "TTS failed" | |
| def check_password(password): | |
| return password == DEVELOPER_PASSWORD | |
| def run_speech_translation(audio, source_lang, target_lang, youtube_url, target_speaker): | |
| transcription, translated_text, audio_url = transcribe_and_speak(audio, source_lang, target_lang, youtube_url, target_speaker) | |
| return transcription, translated_text, audio_url | |
| with gr.Blocks() as demo: | |
| gr.Markdown("# Speech Translation") | |
| # with gr.Tab("User Mode"): | |
| gr.Markdown("Speak into the microphone, upload an audio file, or provide a YouTube URL. The app will translate and speak it back to you.") | |
| with gr.Row(): | |
| user_audio_input = gr.Audio(sources=["microphone", "upload"], type="filepath") | |
| user_youtube_url = gr.Textbox(label="YouTube URL (optional)") | |
| with gr.Row(): | |
| user_source_lang = gr.Dropdown(choices=["en", "ma", "ta", "zh"], label="Source Language", value="en") | |
| user_target_lang = gr.Dropdown(choices=["en", "ma", "ta", "zh"], label="Target Language", value="zh") | |
| user_target_speaker = gr.Dropdown(choices=AVAILABLE_SPEAKERS['zh'], label="Target Speaker", value="childChinese2") | |
| with gr.Row(): | |
| user_button = gr.Button("Translate and Speak", interactive=False) | |
| with gr.Row(): | |
| user_transcription_output = gr.Textbox(label="Transcription") | |
| user_translation_output = gr.Textbox(label="Translation") | |
| user_audio_output = gr.Audio(label="Translated Speech") | |
| user_video_output = gr.HTML(label="YouTube Video") | |
| def update_button_state(audio, youtube_url): | |
| print(audio, youtube_url) | |
| return gr.Button(interactive=bool(audio) or bool(youtube_url)) | |
| user_audio_input.change( | |
| fn=update_button_state, | |
| inputs=[user_audio_input, user_youtube_url], | |
| outputs=user_button | |
| ) | |
| user_youtube_url.change( | |
| fn=update_button_state, | |
| inputs=[user_audio_input, user_youtube_url], | |
| outputs=user_button | |
| ) | |
| user_button.click( | |
| fn=run_speech_translation, | |
| inputs=[user_audio_input, user_source_lang, user_target_lang, user_youtube_url, user_target_speaker], | |
| outputs=[user_transcription_output, user_translation_output, user_audio_output] | |
| ) | |
| def update_video_embed(youtube_url): | |
| if youtube_url: | |
| try: | |
| video_id = fetch_youtube_id(youtube_url) | |
| return f'<iframe width="560" height="315" src="https://www.youtube.com/embed/{video_id}" frameborder="0" allow="autoplay; encrypted-media" allowfullscreen></iframe>' | |
| except Exception as e: | |
| print(f"Error embedding video: {e}") | |
| return "" | |
| user_youtube_url.change( | |
| fn=update_video_embed, | |
| inputs=[user_youtube_url], | |
| outputs=[user_video_output] | |
| ) | |
| def update_target_speakers(target_lang): | |
| return gr.Dropdown(choices=AVAILABLE_SPEAKERS[target_lang], value=AVAILABLE_SPEAKERS[target_lang][0]) | |
| user_target_lang.change( | |
| fn=update_target_speakers, | |
| inputs=[user_target_lang], | |
| outputs=[user_target_speaker] | |
| ) | |
| demo.launch(auth=(os.getenv("DEV_USER"), os.getenv("DEV_PWD"))) |