Update app.py

app.py

@@ -8,22 +8,6 @@ import edge_tts
 import asyncio
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 
-MAX_MAX_NEW_TOKENS = 2048
-DEFAULT_MAX_NEW_TOKENS = 1024
-MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-model_id = "prithivMLmods/FastThink-0.5B-Tiny"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
-)
-model.eval()
-
-
 DESCRIPTION = """
 # QwQ Edge 💬
 """
@@ -42,25 +26,38 @@ h1 {
 }
 '''
 
-# List of voices
-voices = [
+MAX_MAX_NEW_TOKENS = 2048
+DEFAULT_MAX_NEW_TOKENS = 1024
+MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+model_id = "prithivMLmods/FastThink-0.5B-Tiny"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+)
+model.eval()
+
+TTS_VOICES = [
     "en-US-JennyNeural",  # @tts1
     "en-US-GuyNeural",    # @tts2
     "en-US-AriaNeural",   # @tts3
-    "en-US-JaneNeural",   # @tts4
-    "en-US-JasonNeural",  # @tts5
-    "en-US-NancyNeural",  # @tts6
-    "en-US-TonyNeural",   # @tts7
+    "en-US-DavisNeural",  # @tts4
+    "en-US-JaneNeural",   # @tts5
+    "en-US-JasonNeural",  # @tts6
+    "en-US-NancyNeural",  # @tts7
+    "en-US-TonyNeural",   # @tts8
 ]
 
-
 async def text_to_speech(text: str, voice: str, output_file="output.mp3"):
     """Convert text to speech using Edge TTS and save as MP3"""
     communicate = edge_tts.Communicate(text, voice)
     await communicate.save(output_file)
     return output_file
 
-
 @spaces.GPU
 def generate(
     message: str,
@@ -72,25 +69,16 @@ def generate(
     repetition_penalty: float = 1.2,
 ):
     """Generates chatbot response and handles TTS requests"""
-    is_tts = message.strip().lower().startswith("@tts")
-    tts_index = None
-
-    if is_tts:
-        # Extract the number after @tts
-        tts_part = message.strip().lower().split()[0]  # Get the @ttsX part
-        if len(tts_part) > 8:  # Check if it's @ttsX (e.g., @tts1, @tts2, etc.)
-            try:
-                tts_index = int(tts_part[8:]) - 1  # Convert to 0-based index
-                if tts_index < 0 or tts_index >= len(voices):
-                    gr.Warning(f"Invalid TTS voice index. Using default voice.")
-                    tts_index = 0
-            except ValueError:
-                gr.Warning(f"Invalid TTS voice index. Using default voice.")
-                tts_index = 0
-        else:
-            tts_index = 0  # Default to the first voice if no number is provided
-
-        message = message.replace(tts_part, "").strip()  # Remove @ttsX from the message
+    tts_prefix = "@tts"
+    is_tts = any(message.strip().lower().startswith(f"{tts_prefix}{i}") for i in range(1, 9))
+    voice_index = next((i for i in range(1, 9) if message.strip().lower().startswith(f"{tts_prefix}{i}")), None)
+    
+    if is_tts and voice_index:
+        voice = TTS_VOICES[voice_index - 1]
+        message = message.replace(f"{tts_prefix}{voice_index}", "").strip()
+    else:
+        voice = None
+        message = message.replace(tts_prefix, "").strip()
 
     conversation = [*chat_history, {"role": "user", "content": message}]
 
@@ -122,14 +110,12 @@ def generate(
 
     final_response = "".join(outputs)
 
-    if is_tts:
-        voice = voices[tts_index]  # Select the voice based on the index
+    if is_tts and voice:
         output_file = asyncio.run(text_to_speech(final_response, voice))
         yield gr.Audio(output_file, autoplay=True)  # Return playable audio
     else:
         yield final_response  # Return text response
 
-
 demo = gr.ChatInterface(
     fn=generate,
     additional_inputs=[
@@ -141,12 +127,12 @@ demo = gr.ChatInterface(
     ],
     stop_btn=None,
     examples=[
-        ["@tts7 Who is Nikola Tesla, and why did he die?"],
+        ["@tts1 Who is Nikola Tesla, and why did he die?"],
         ["A train travels 60 kilometers per hour. If it travels for 5 hours, how far will it travel in total?"],
         ["Write a Python function to check if a number is prime."],
-        ["@tts6 What causes rainbows to form?"],
+        ["@tts2 What causes rainbows to form?"],
         ["Rewrite the following sentence in passive voice: 'The dog chased the cat.'"],
-        ["@tts4 What is the capital of France?"],
+        ["@tts5 What is the capital of France?"],
     ],
     cache_examples=False,
     type="messages",