app.py

import gradio as gr
import torch
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
from IndicTransToolkit.processor import IndicProcessor
import requests
from datetime import datetime
import tempfile
from gtts import gTTS
import os

# Device configuration
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# Supabase configuration
SUPABASE_URL = "https://gptmdbhzblfybdnohqnh.supabase.co"
SUPABASE_API_KEY = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9..."

# Load translation models
model_en_to_indic = AutoModelForSeq2SeqLM.from_pretrained("ai4bharat/indictrans2-en-indic-1B", trust_remote_code=True).to(DEVICE)
tokenizer_en_to_indic = AutoTokenizer.from_pretrained("ai4bharat/indictrans2-en-indic-1B", trust_remote_code=True)
model_indic_to_en = AutoModelForSeq2SeqLM.from_pretrained("ai4bharat/indictrans2-indic-en-1B", trust_remote_code=True).to(DEVICE)
tokenizer_indic_to_en = AutoTokenizer.from_pretrained("ai4bharat/indictrans2-indic-en-1B", trust_remote_code=True)

ip = IndicProcessor(inference=True)
asr = pipeline("automatic-speech-recognition", model="openai/whisper-small")


def save_to_supabase(input_text, output_text, direction):
    if not input_text.strip() or not output_text.strip():
        return "Nothing to save."

    table_name = "translations" if direction == "en_to_ks" else "ks_to_en_translations"
    payload = {
        "timestamp": datetime.utcnow().isoformat(),
        "input_text": input_text,
        "output_text": output_text
    }

    headers = {
        "apikey": SUPABASE_API_KEY,
        "Authorization": f"Bearer {SUPABASE_API_KEY}",
        "Content-Type": "application/json"
    }

    try:
        response = requests.post(
            f"{SUPABASE_URL}/rest/v1/{table_name}",
            headers=headers,
            json=payload,
            timeout=10
        )
        return "✅ Saved successfully!" if response.status_code == 201 else "❌ Failed to save."
    except Exception as e:
        print("SAVE EXCEPTION:", e)
        return "❌ Save request error."


def get_translation_history(direction="en_to_ks"):
    table_name = "translations" if direction == "en_to_ks" else "ks_to_en_translations"

    headers = {
        "apikey": SUPABASE_API_KEY,
        "Authorization": f"Bearer {SUPABASE_API_KEY}"
    }

    try:
        response = requests.get(
            f"{SUPABASE_URL}/rest/v1/{table_name}?order=timestamp.desc&limit=10",
            headers=headers,
            timeout=10
        )
        if response.status_code == 200:
            records = response.json()
            return "\n\n".join(
                [f"Input: {r['input_text']} → Output: {r['output_text']}" for r in records]
            )
        return "Failed to load history."
    except Exception as e:
        print("HISTORY FETCH ERROR:", e)
        return "Error loading history."


def translate(text, direction):
    if not text.strip():
        return "Please enter some text.", gr.update(), gr.update()

    if direction == "en_to_ks":
        src_lang = "eng_Latn"
        tgt_lang = "kas_Arab"
        model = model_en_to_indic
        tokenizer = tokenizer_en_to_indic
    else:
        src_lang = "kas_Arab"
        tgt_lang = "eng_Latn"
        model = model_indic_to_en
        tokenizer = tokenizer_indic_to_en

    try:
        processed = ip.preprocess_batch([text], src_lang=src_lang, tgt_lang=tgt_lang)
        batch = tokenizer(processed, return_tensors="pt", padding=True).to(DEVICE)

        with torch.no_grad():
            outputs = model.generate(
                **batch,
                max_length=256,
                num_beams=5,
                num_return_sequences=1
            )

        translated = tokenizer.batch_decode(outputs, skip_special_tokens=True)
        result = ip.postprocess_batch(translated, lang=tgt_lang)[0]

        return result, gr.update(), gr.update()
    except Exception as e:
        print("Translation Error:", e)
        return "⚠️ Translation failed.", gr.update(), gr.update()


def transcribe_audio(audio_path):
    try:
        result = asr(audio_path)
        return result["text"]
    except Exception as e:
        print("STT Error:", e)
        return "⚠️ Transcription failed."


def synthesize_tts(text, direction):
    if direction == "ks_to_en" and text.strip():
        try:
            tts = gTTS(text=text, lang="en")
            tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
            tts.save(tmp_file.name)
            tmp_file.close()
            return tmp_file.name
        except Exception as e:
            print("TTS Error:", e)
    return None


def switch_direction(direction, input_text_val, output_text_val):
    new_direction = "ks_to_en" if direction == "en_to_ks" else "en_to_ks"
    input_label = "Kashmiri Text" if new_direction == "ks_to_en" else "English Text"
    output_label = "English Translation" if new_direction == "ks_to_en" else "Kashmiri Translation"
    return (
        new_direction,
        gr.update(value=output_text_val, label=input_label),
        gr.update(value=input_text_val, label=output_label)
    )


def handle_audio_translation(audio_path, direction):
    if direction == "en_to_ks":
        transcription = transcribe_audio(audio_path)
        if transcription.startswith("⚠️"):
            return transcription, "", "", None
        translation, _, _ = translate(transcription, direction)
        return transcription, translation, transcription, None
    else:
        # Assume audio_path is not used; rely on text in the input box
        transcription = transcribe_audio(audio_path)
        translation, _, _ = translate(transcription, direction)
        tts_audio = synthesize_tts(translation, direction)
        return transcription, translation, transcription, tts_audio


# === Gradio Interface ===
with gr.Blocks() as interface:
    gr.HTML("""
    <div style="display: flex; justify-content: space-between; align-items: center; padding: 10px;">
        <img src="https://raw.githubusercontent.com/BurhaanRasheedZargar/Images/211321a234613a9c3dd944fe9367cf13d1386239/assets/left_logo.png" style="height:150px; width:auto;">
        <h2 style="margin: 0; text-align: center;">English ↔ Kashmiri Translator</h2>
        <img src="https://raw.githubusercontent.com/BurhaanRasheedZargar/Images/77797f7f7cbee328fa0f9d31cf3e290441e04cd3/assets/right_logo.png">
    </div>
    """)

    translation_direction = gr.State(value="en_to_ks")

    with gr.Row():
        input_text = gr.Textbox(lines=2, label="English Text", placeholder="Enter text....")
        output_text = gr.Textbox(lines=2, label="Kashmiri Translation", placeholder="Translated text....")

    with gr.Row():
        translate_button = gr.Button("Translate")
        save_button = gr.Button("Save Translation")
        switch_button = gr.Button("Switch")

    save_status = gr.Textbox(label="Save Status", interactive=False)
    history_box = gr.Textbox(lines=10, label="Translation History", interactive=False)

    with gr.Row():
        audio_input = gr.Audio(type="filepath", label="🎙️ Upload or record English audio")
        audio_output = gr.Audio(label="🔊 English Output Audio")

    stt_translate_button = gr.Button("🎤 Transcribe & Translate")

    # Events
    translate_button.click(
        fn=translate,
        inputs=[input_text, translation_direction],
        outputs=[output_text, input_text, output_text]
    )

    save_button.click(
        fn=save_to_supabase,
        inputs=[input_text, output_text, translation_direction],
        outputs=save_status
    ).then(
        fn=get_translation_history,
        inputs=translation_direction,
        outputs=history_box
    )

    switch_button.click(
        fn=switch_direction,
        inputs=[translation_direction, input_text, output_text],
        outputs=[translation_direction, input_text, output_text]
    )

    stt_translate_button.click(
        fn=handle_audio_translation,
        inputs=[audio_input, translation_direction],
        outputs=[input_text, output_text, input_text, audio_output]
    )

if __name__ == "__main__":
    interface.queue().launch(share=True)