Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,339 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import os
+import psutil
+import signal
+import time
+import re
+import requests
+import subprocess
 import streamlit as st
+from fastapi import FastAPI, File, UploadFile, Form
+from fastapi.responses import JSONResponse
+from faster_whisper import WhisperModel
+from pydub import AudioSegment
+from io import BytesIO
+import torch
+import logging
+import uvicorn
+from pyngrok import ngrok
+import nest_asyncio
+from tenacity import retry, stop_after_attempt, wait_fixed
+from deep_translator import GoogleTranslator
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# GitHub Gist configuration
+GITHUB_TOKEN = "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"  # Replace with your GitHub PAT
+GIST_ID = "e1b9e5a9d7167ab2458405c38a2803c2"
+GIST_FILENAME = "ngrok_public_url.txt"
+GIST_DESCRIPTION = "Ngrok Public URL for FastAPI Server"
+
+# Function to free port 8000
+def free_port(port=8000):
+    try:
+        for proc in psutil.process_iter(['pid', 'name']):
+            for conn in proc.net_connections(kind='inet'):
+                if conn.laddr.port == port:
+                    logger.info(f"Terminating process {proc.pid} using port {port}")
+                    os.kill(proc.pid, signal.SIGTERM)
+                    proc.wait(timeout=5)
+    except Exception as e:
+        logger.warning(f"Error freeing port {port}: {str(e)}")
+
+# Function to create or update a GitHub Gist
+@retry(stop=stop_after_attempt(3), wait=wait_fixed(2))
+def update_github_gist(public_url: str) -> str:
+    global GIST_ID
+    headers = {
+        "Authorization": f"token {GITHUB_TOKEN}",
+        "Accept": "application/vnd.github.v3+json"
+    }
+    gist_data = {
+        "description": GIST_DESCRIPTION,
+        "public": False,
+        "files": {
+            GIST_FILENAME: {
+                "content": public_url
+            }
+        }
+    }
+
+    try:
+        if GIST_ID:
+            response = requests.patch(
+                f"https://api.github.com/gists/{GIST_ID}",
+                headers=headers,
+                json=gist_data
+            )
+            if response.status_code == 200:
+                logger.info(f"Successfully updated Gist {GIST_ID} with URL: {public_url}")
+                return GIST_ID
+            else:
+                logger.error(f"Failed to update Gist {GIST_ID}: {response.text}")
+                raise Exception(f"Failed to update Gist: {response.text}")
+        else:
+            response = requests.post(
+                "https://api.github.com/gists",
+                headers=headers,
+                json=gist_data
+            )
+            if response.status_code == 201:
+                GIST_ID = response.json()['id']
+                logger.info(f"Created new Gist {GIST_ID} with URL: {public_url}")
+                return GIST_ID
+            else:
+                logger.error(f"Failed to create Gist: {response.text}")
+                raise Exception(f"Failed to create Gist: {response.text}")
+    except Exception as e:
+        logger.error(f"Error updating/creating Gist: {str(e)}")
+        raise
+
+# Function to start ngrok tunnel
+def start_ngrok_tunnel(port=8000):
+    try:
+        tunnel = ngrok.connect(port, bind_tls=True)
+        public_url = tunnel.public_url
+        logger.info(f"ngrok tunnel created at {public_url}")
+        return public_url
+    except Exception as e:
+        logger.error(f"Failed to start ngrok tunnel: {str(e)}")
+        raise
+
+# Initialize FastAPI app
+app = FastAPI()
+
+# Load Whisper model
+model = None
+try:
+    model = WhisperModel("large", device="cuda", compute_type="float16")
+    logger.info("Whisper model loaded on GPU")
+except Exception as e:
+    logger.warning(f"Failed to load Whisper model on GPU: {str(e)}. Falling back to CPU.")
+    try:
+        model = WhisperModel("large", device="cpu", compute_type="int8")
+        logger.info("Whisper model loaded on CPU")
+    except Exception as cpu_e:
+        logger.error(f"Failed to load Whisper model on CPU: {str(cpu_e)}")
+        raise RuntimeError(f"Failed to load Whisper model on both GPU and CPU: {str(cpu_e)}")
+
+def clean_transcription(text):
+    """Clean the transcription by removing consecutive duplicates"""
+    sentences = [s.strip() for s in re.split(r'[.!?]', text) if s.strip()]
+    unique_sentences = []
+    prev_sentence = None
+    
+    for sentence in sentences:
+        if sentence != prev_sentence:
+            unique_sentences.append(sentence)
+            prev_sentence = sentence
+    
+    processed_sentences = []
+    for sentence in unique_sentences:
+        words = sentence.split()
+        unique_words = []
+        prev_word = None
+        for word in words:
+            if word != prev_word:
+                unique_words.append(word)
+                prev_word = word
+        processed_sentences.append(' '.join(unique_words))
+    
+    cleaned_text = '. '.join(processed_sentences)
+    if text.endswith('.'):
+        cleaned_text += '.'
+    elif text.endswith('?'):
+        cleaned_text += '?'
+    elif text.endswith('!'):
+        cleaned_text += '!'
+    
+    return cleaned_text
+
+def translate_to_language(text_or_dict, source_lang, target_lang):
+    lang_codes = {
+        'en': 'en', 'hi': 'hi', 'bn': 'bn', 'te': 'te', 'mr': 'mr',
+        'ta': 'ta', 'ur': 'ur', 'gu': 'gu', 'kn': 'kn', 'ml': 'ml',
+        'pa': 'pa', 'or': 'or', 'as': 'as', 'ne': 'ne'
+    }
+    malayalam_reverse_map = {
+        'ആശയവുമില്ല': 'ഐഡിയയുമില്ല',
+        'എന്ന് ': 'നിന്നും ',
+        'ചെയ്യാൻ': 'ചെയ്യാവോ',
+        # ... (keeping the full malayalam_reverse_map as in the original code)
+        'നികുതി': 'ടാക്സ്'
+    }
+    try:
+        if isinstance(text_or_dict, dict):
+            text = text_or_dict.get('raw_transcription', '')
+            if not text:
+                logger.warning("No raw_transcription found in input dictionary")
+                return {'translated_text': ''}
+        else:
+            text = text_or_dict
+            if not isinstance(text, str):
+                logger.warning(f"Input text_or_dict is not a string or dict: {type(text_or_dict)}")
+                text = str(text_or_dict)
+        text = text.strip()
+        if not text:
+            logger.warning("Input text is empty after stripping")
+            return {'translated_text': ''}
+        text = re.sub(r'[^\w\s.,!?]', '', text)
+        if source_lang not in lang_codes or target_lang not in lang_codes:
+            logger.error(f"Unsupported language: source={source_lang}, target={target_lang}")
+            return {'translated_text': '', 'error': f"Unsupported language: source={source_lang}, target={target_lang}"}
+        source_code = lang_codes.get(source_lang, 'en')
+        target_code = lang_codes.get(target_lang, 'ml')
+        logger.info(f"Translating from {source_lang} ({source_code}) to {target_lang} ({target_code})")
+        translator = GoogleTranslator(source=source_code, target=target_code)
+        translated_text = translator.translate(text)
+        if not translated_text:
+            logger.warning("Translation returned empty result")
+            return {'translated_text': ''}
+        if target_lang == 'ml':
+            for incorrect, correct in malayalam_reverse_map.items():
+                translated_text = translated_text.replace(incorrect, correct)
+        logger.debug(f"Translated text: {translated_text[:100]}...")
+        return {'translated_text': translated_text}
+    except Exception as e:
+        logger.error(f"Translation failed: {str(e)}")
+        return {'translated_text': '', 'error': f"Translation failed: {str(e)}"}
+
+@app.post("/transcribe")
+async def transcribe(
+    audio: UploadFile = File(...),
+    model_size: str = Form("small"),
+    transcription_language: str = Form("en"),
+    target_language: str = Form(None)
+):
+    try:
+        audio_data = await audio.read()
+        audio_buffer = BytesIO(audio_data)
+        audio_segment = AudioSegment.from_file(audio_buffer, format="wav")
+        wav_buffer = BytesIO()
+        audio_segment.export(wav_buffer, format="wav")
+        wav_buffer.seek(0)
+        segments, info = model.transcribe(
+            wav_buffer,
+            language=transcription_language,
+            task="transcribe",
+            beam_size=1,
+            vad_filter=True
+        )
+        
+        raw_segments = [segment.text for segment in segments]
+        raw_transcription = " ".join(raw_segments)
+        raw_transcription = clean_transcription(raw_transcription)
+        
+        response = {
+            "raw_transcription": raw_transcription,
+            "audio_metadata": {
+                "extension": audio.filename.split('.')[-1],
+                "duration": len(audio_segment) / 1000,
+                "emotion": "unknown"
+            },
+            "detected_language": info.language,
+            "detected_language_name": {
+                'en': 'English', 'hi': 'Hindi', 'bn': 'Bengali', 'te': 'Telugu', 'mr': 'Marathi',
+                'ta': 'Tamil', 'ur': 'Urdu', 'gu': 'Gujarati', 'kn': 'Kannada', 'ml': 'Malayalam',
+                'pa': 'Punjabi', 'or': 'Odia', 'as': 'Assamese', 'ne': 'Nepali'
+            }.get(info.language, info.language)
+        }
+        
+        if target_language:
+            translation_result = translate_to_language(
+                {'raw_transcription': raw_transcription},
+                source_lang=info.language,
+                target_lang=target_language
+            )
+            response['translated_text'] = translation_result.get('translated_text', '')
+            if 'error' in translation_result:
+                response['translation_error'] = translation_result['error']
+                
+        return JSONResponse(response)
+    except Exception as e:
+        logger.error(f"Transcription failed: {str(e)}")
+        return JSONResponse({"error": str(e)}, status_code=500)
+
+# Function to start the FastAPI server as a subprocess
+def start_server():
+    try:
+        free_port(8000)
+        nest_asyncio.apply()
+        ngrok.set_auth_token("2zXsNdquBfJFwENcejztPfNEdan_57vHEJ6RhnFxNoaNxR5cW")  # Replace with your token
+        public_url = start_ngrok_tunnel(8000)
+        st.session_state['public_url'] = public_url
+        st.session_state['gist_id'] = update_github_gist(str(public_url))
+        
+        # Start the FastAPI server as a subprocess
+        process = subprocess.Popen(
+            ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True
+        )
+        st.session_state['server_process'] = process
+        st.success(f"Server started successfully! Public URL: {public_url}")
+        st.write(f"Updated Gist ID: {st.session_state['gist_id']}")
+        return process
+    except Exception as e:
+        logger.error(f"Failed to start server: {str(e)}")
+        st.error("Error: Could not start server. Please check the ngrok dashboard (https://dashboard.ngrok.com/agents) to terminate existing sessions or upgrade your ngrok account.")
+        return None
+
+# Function to stop the FastAPI server and ngrok tunnel
+def stop_server():
+    try:
+        # Terminate the server process
+        if 'server_process' in st.session_state and st.session_state['server_process']:
+            process = st.session_state['server_process']
+            process.terminate()
+            process.wait(timeout=5)
+            logger.info("Server process terminated")
+        
+        # Terminate ngrok tunnel
+        ngrok.kill()
+        logger.info("ngrok tunnel terminated")
+        
+        # Clear session state
+        st.session_state['server_process'] = None
+        st.session_state['public_url'] = None
+        st.session_state['gist_id'] = None
+        st.success("Server and ngrok tunnel stopped successfully")
+    except Exception as e:
+        logger.error(f"Failed to stop server: {str(e)}")
+        st.error(f"Error stopping server: {str(e)}")
+
+# Streamlit app
+def main():
+    st.title("FastAPI Server Control")
+    st.write("Use the buttons below to start or stop the FastAPI server and ngrok tunnel.")
+
+    # Initialize session state
+    if 'server_process' not in st.session_state:
+        st.session_state['server_process'] = None
+        st.session_state['public_url'] = None
+        st.session_state['gist_id'] = None
+
+    # Start Server Button
+    if st.button("Start Server"):
+        if st.session_state['server_process'] is None:
+            start_server()
+        else:
+            st.warning("Server is already running!")
+
+    # Stop Server Button
+    if st.button("Stop Server"):
+        if st.session_state['server_process'] is not None:
+            stop_server()
+        else:
+            st.warning("No server is running!")
+
+    # Display current status
+    if st.session_state['public_url']:
+        st.write(f"**Current Public URL**: {st.session_state['public_url']}")
+        st.write(f"**Gist ID**: {st.session_state['gist_id']}")
+    else:
+        st.write("**Status**: Server is not running")
+
+if __name__ == "__main__":
+    main()