app.py

#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 = "cuda" if torch.cuda.is_available() else "cpu"

# Load 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")

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

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

    table = "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}", json=payload, headers=headers)
        return "✅ Saved successfully!" if response.status_code == 201 else "❌ Failed to save."
    except Exception as e:
        print("Save error:", e)
        return "❌ Save error."

def get_translation_history(direction):
    table = "translations" if direction == "en_to_ks" else "ks_to_en_translations"
    headers = {
        "apikey": SUPABASE_API_KEY,
        "Authorization": f"Bearer {SUPABASE_API_KEY}"
    }

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

# --- Translation ---
def translate(text, direction):
    if not text.strip():
        return "Enter some text.", None

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

    try:
        batch = ip.preprocess_batch([text], src_lang=src_lang, tgt_lang=tgt_lang)
        tokens = tokenizer(batch, return_tensors="pt", padding=True).to(DEVICE)
        with torch.no_grad():
            output = model.generate(**tokens, max_length=256, num_beams=5)
        result = tokenizer.batch_decode(output, skip_special_tokens=True)
        final = ip.postprocess_batch(result, lang=tgt_lang)[0]
        return final, None
    except Exception as e:
        print("Translation error:", e)
        return "⚠️ Translation failed.", None

# --- TTS for English output ---
def synthesize_tts(text):
    try:
        tts = gTTS(text=text, lang="en")
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as f:
            tts.save(f.name)
            return f.name
    except Exception as e:
        print("TTS error:", e)
        return None

# --- STT only for en_to_ks ---
def generate_stt_for_input(audio_path, direction):
    if direction != "en_to_ks":
        return "⚠️ Audio input is only supported for English to Kashmiri.", "", None

    try:
        transcription = asr(audio_path)["text"]
    except Exception as e:
        print("STT error:", e)
        return "⚠️ Transcription failed.", "", None

    translated, _ = translate(transcription, direction)
    return transcription, translated, None

# --- Generate TTS for English output ---
def generate_tts_for_output(output_text, direction):
    if direction == "ks_to_en" and output_text.strip():
        return synthesize_tts(output_text)
    return None

# --- Switch UI direction ---
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), None

# === Gradio Interface ===
with gr.Blocks() as interface:
    gr.HTML("<h2 style='text-align:center;'>English ↔ Kashmiri Translator</h2>")
    translation_direction = gr.State(value="en_to_ks")

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

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

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

    with gr.Row():
        audio_input = gr.Audio(type="filepath", label="🎙️ Record English audio")
        audio_output = gr.Audio(label="🔊 English TTS", interactive=False)

    stt_button = gr.Button("🎤 Transcribe & Translate (EN → KS Only)")
    tts_button = gr.Button("🔊 Generate English Speech (KS → EN Only)")

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

    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
    )

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

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

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