Create app.py

app.py

@@ -0,0 +1,198 @@
+import spaces
+import gradio as gr
+import torch
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    pipeline,
+    AutoProcessor,
+    AutoModelForSpeechSeq2Seq,
+    BitsAndBytesConfig
+)
+from datasets import load_dataset
+import numpy as np
+from transformers import AutoModelForTextToSpeech, SpeechT5HifiGan
+import torchaudio
+
+@spaces.GPU
+def dummy(): # just a dummy
+    pass
+
+# Constants
+# DEVICE = "cpu"
+LANGUAGE_CODES = {
+    "English": "en",
+    "Chinese": "zh"
+}
+
+# Initialize components with efficient settings
+def initialize_components():
+    # Use XVERSE-13B-Chat as the base model - good multilingual support and reasonable size
+    # Load in 4-bit quantization to reduce memory usage
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_compute_dtype=torch.float16,
+    )
+    
+    llm = AutoModelForCausalLM.from_pretrained(
+        "xverse/XVERSE-13B-Chat",
+        quantization_config=bnb_config,
+        device_map="auto"
+    )
+    tokenizer = AutoTokenizer.from_pretrained("xverse/XVERSE-13B-Chat")
+
+    # Whisper model for STT (small for efficiency)
+    processor = AutoProcessor.from_pretrained("openai/whisper-small")
+    stt_model = AutoModelForSpeechSeq2Seq.from_pretrained(
+        "openai/whisper-small",
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True,
+    )
+
+    # VITS for TTS (supports both English and Chinese)
+    tts_model = load_model("facebook/mms-tts-eng)
+    vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
+
+    return llm, tokenizer, processor, stt_model, tts_model, vocoder
+
+def load_model(model_name):
+    """Helper function to load models with optimized settings"""
+    return AutoModelForTextToSpeech.from_pretrained(
+        model_name,
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True,
+    )
+
+class ConversationManager:
+    def __init__(self):
+        self.history = []
+    
+    def add_message(self, role, content, audio_path=None):
+        self.history.append({
+            "role": role,
+            "content": content,
+            "audio_path": audio_path
+        })
+    
+    def get_formatted_history(self):
+        return "\n".join([
+            f"{msg['role']}: {msg['content']}" for msg in self.history
+        ])
+
+def speech_to_text(audio, processor, model, target_language):
+    """Convert speech to text using Whisper"""
+    input_features = processor(
+        audio, 
+        sampling_rate=16000, 
+        return_tensors="pt"
+    ).input_features
+    
+    predicted_ids = model.generate(
+        input_features,
+        language=LANGUAGE_CODES[target_language]
+    )
+    
+    transcription = processor.batch_decode(
+        predicted_ids, 
+        skip_special_tokens=True
+    )[0]
+    return transcription
+
+def generate_response(prompt, llm, tokenizer):
+    """Generate LLM response with optimized settings"""
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = llm.generate(
+        **inputs,
+        max_length=512,
+        num_return_sequences=1,
+        temperature=0.7,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
+    )
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+
+def text_to_speech(text, model, vocoder, language):
+    """Convert text to speech using MMS-TTS"""
+    inputs = processor(text, return_tensors="pt")
+    speech = model.generate_speech(inputs["input_ids"], vocoder)
+    return speech
+
+def create_gradio_interface():
+    # Initialize components
+    llm, tokenizer, processor, stt_model, tts_model, vocoder = initialize_components()
+    conversation_manager = ConversationManager()
+    
+    with gr.Blocks() as interface:
+        with gr.Row():
+            language_selector = gr.Dropdown(
+                choices=list(LANGUAGE_CODES.keys()),
+                value="English",
+                label="Select Language"
+            )
+        
+        with gr.Row():
+            # Audio input
+            audio_input = gr.Audio(
+                source="microphone",
+                type="numpy",
+                label="Speak"
+            )
+        
+        with gr.Row():
+            # Chat history display
+            chat_display = gr.Textbox(
+                value="",
+                label="Conversation History",
+                lines=10,
+                readonly=True
+            )
+        
+        with gr.Row():
+            # Assistant's audio response
+            audio_output = gr.Audio(
+                label="Assistant's Response",
+                type="numpy"
+            )
+
+        def process_conversation(audio, language):
+            # Speech to text
+            user_text = speech_to_text(
+                audio,
+                processor,
+                stt_model,
+                language
+            )
+            conversation_manager.add_message("User", user_text)
+            
+            # Generate LLM response
+            context = conversation_manager.get_formatted_history()
+            response = generate_response(context, llm, tokenizer)
+            conversation_manager.add_message("Assistant", response)
+            
+            # Text to speech
+            speech_output = text_to_speech(
+                response,
+                tts_model,
+                vocoder,
+                language
+            )
+            
+            return (
+                conversation_manager.get_formatted_history(),
+                (16000, speech_output.numpy())
+            )
+
+        audio_input.change(
+            process_conversation,
+            inputs=[audio_input, language_selector],
+            outputs=[chat_display, audio_output]
+        )
+    
+    return interface
+
+# Launch the application
+if __name__ == "__main__":
+    interface = create_gradio_interface()
+    interface.launch()