app.py

import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import os
import json
import gspread
from oauth2client.service_account import ServiceAccountCredentials
from datetime import datetime
from gtts import gTTS
import tempfile
import requests

# --- CONFIGURATION ---
MODEL_K2H_REPO = "ankitklakra/kurukh-to-hindi"
MODEL_H2K_REPO = "ankitklakra/hindi-to-kurukh"
SHEET_NAME = "Kurukh_Feedback_Log"

# --- LOAD MODELS ---
print("Loading Translation Models...")
tokenizer = AutoTokenizer.from_pretrained("google/mt5-small")
model_k2h = AutoModelForSeq2SeqLM.from_pretrained(MODEL_K2H_REPO)
model_h2k = AutoModelForSeq2SeqLM.from_pretrained(MODEL_H2K_REPO)

pipe_k2h = pipeline("text2text-generation", model=model_k2h, tokenizer=tokenizer)
pipe_h2k = pipeline("text2text-generation", model=model_h2k, tokenizer=tokenizer)

print("Loading Voice Model...")
asr_pipeline = pipeline("automatic-speech-recognition", model="openai/whisper-tiny")

# --- HELPER FUNCTIONS ---
def transliterate_to_hindi(text):
    try:
        url = "https://inputtools.google.com/request?text={}&itc=hi-t-i0-und&num=1"
        response = requests.get(url.format(text))
        result = response.json()
        return result[1][0][1][0]
    except:
        return text

def save_to_sheet(original, translation, correction, direction):
    try:
        json_creds = os.getenv("GOOGLE_CREDENTIALS")
        if not json_creds: return "⚠️ Error: Credentials missing."
        creds_dict = json.loads(json_creds)
        scope = ["https://spreadsheets.google.com/feeds", "https://www.googleapis.com/auth/drive"]
        creds = ServiceAccountCredentials.from_json_keyfile_dict(creds_dict, scope)
        client = gspread.authorize(creds)
        sheet = client.open(SHEET_NAME).sheet1
        if not sheet.get_all_values():
            sheet.append_row(["Timestamp", "Direction", "Original Text", "AI Translation", "User Correction"])
        sheet.append_row([str(datetime.now()), direction, original, translation, correction])
        return "✅ Saved to Google Sheets."
    except Exception as e:
        return f"❌ Error: {str(e)}"

def speech_to_text(audio_path):
    if audio_path is None: return ""
    return asr_pipeline(audio_path)["text"]

def text_to_speech(text, language="hi"):
    if not text: return None
    try:
        tts = gTTS(text=text, lang=language)
        temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
        tts.save(temp_file.name)
        return temp_file.name
    except: return None

# --- MAIN TRANSLATION LOGIC ---
def process_translation(text, audio_input, direction, is_hinglish):
    original_text = speech_to_text(audio_input) if audio_input else text
    if not original_text: return "", "", None

    if direction == "Hindi -> Kurukh" and is_hinglish:
        original_text = transliterate_to_hindi(original_text)

    target_pipeline = pipe_k2h if direction == "Kurukh -> Hindi" else pipe_h2k
    try:
        results = target_pipeline(
            original_text,
            max_length=128,
            num_beams=5,
            no_repeat_ngram_size=2,
            repetition_penalty=2.0,
            early_stopping=True
        )
        translated_text = results[0]['generated_text']
    except Exception as e:
        return str(e), "", None

    audio_output = None
    if direction == "Kurukh -> Hindi":
        audio_output = text_to_speech(translated_text, "hi")
    
    return original_text, translated_text, audio_output

# --- CUSTOM CSS  ---
custom_css = """
<style>
@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;600&display=swap');
body, button, input, select, textarea, .gradio-container { font-family: 'Poppins', sans-serif !important; }
.header-div { text-align: center; margin-bottom: 20px; }
.header-title { font-size: 2.5em; font-weight: 600; color: #2c3e50; margin: 0; }
.header-subtitle { font-size: 1.2em; color: #7f8c8d; font-weight: 300; margin-top: 5px; }
</style>
"""

# --- THE UI ---

with gr.Blocks(theme=gr.themes.Soft(), title="Kurukh AI Translator") as demo:
    
    # INJECT CSS
    gr.HTML(custom_css)

    # HEADER
    gr.HTML("""
    <div class="header-div">
        <h1 class="header-title">🇮🇳 AI Kurukh (Oraon) Translator</h1>
        <p class="header-subtitle">Bridging Communities with Artificial Intelligence | Voice & Hinglish Supported</p>
    </div>
    """)

    with gr.Tabs():
        
        # --- TAB 1: TRANSLATOR ---
        with gr.TabItem("🗣️ Translator"):
            
            with gr.Accordion("ℹ️ How to use (Click to expand)", open=False):
                gr.Markdown("""
                1. **Select Mode:** Kurukh -> Hindi OR Hindi -> Kurukh.
                2. **Hinglish:** Check the box if you type Hindi in English (e.g., 'Tumhara').
                3. **Voice:** Click the microphone to speak (Hindi only).
                """)

            with gr.Row():
                # LEFT COLUMN
                with gr.Column():
                    gr.Markdown("### 📥 Input Source")
                    
                    with gr.Row():
                        direction = gr.Radio(["Kurukh -> Hindi", "Hindi -> Kurukh"], label="Translation Mode", value="Kurukh -> Hindi")
                    
                    is_hinglish = gr.Checkbox(label="🔤 Hinglish Typing (e.g. 'Tumhara')", value=False)
                    
                    input_text = gr.Textbox(label="Enter Text", placeholder="Type sentences here...", lines=4)
                    
                   
                    input_audio = gr.Audio(sources=["microphone"], type="filepath", label="🎙️ Voice Input (Hindi Only)")
                    
                    translate_btn = gr.Button("Translate 🚀", variant="primary")

                # RIGHT COLUMN
                with gr.Column():
                    gr.Markdown("### 📤 Translation Result")
                  
                    output_text = gr.Textbox(label="Translation", lines=4, interactive=False, show_copy_button=True)
                    output_audio = gr.Audio(label="🔊 Listen (Hindi Only)", interactive=False, autoplay=False)

            translate_btn.click(
                fn=process_translation, 
                inputs=[input_text, input_audio, direction, is_hinglish], 
                outputs=[input_text, output_text, output_audio]
            )

        # --- TAB 2: FEEDBACK ---
        with gr.TabItem("📝 Improve the AI"):
            gr.Markdown("### 🛠️ Help us improve accuracy")
            gr.Markdown("Our AI is learning! If you spot a mistake, please submit the correct translation below.")
            
            with gr.Column():
                with gr.Row():
                    fb_direction = gr.Radio(["Kurukh -> Hindi", "Hindi -> Kurukh"], label="Direction", value="Kurukh -> Hindi")
                
                with gr.Row():
                    fb_original = gr.Textbox(label="Original Text")
                    fb_ai_output = gr.Textbox(label="AI's Wrong Translation")
                
                fb_user_correct = gr.Textbox(label="✅ Your Correct Translation", lines=2, placeholder="Type the correct translation here")
                
                submit_btn = gr.Button("Submit Correction to Database", variant="secondary")
                status_lbl = gr.Label(label="Status")
                
                submit_btn.click(
                    fn=save_to_sheet, 
                    inputs=[fb_original, fb_ai_output, fb_user_correct, fb_direction], 
                    outputs=status_lbl
                )
    
    gr.Markdown("---")
    gr.HTML("<center style='color: #888;'>Built with ❤️ for the Kurukh Community • Powered by Google mT5 & OpenAI Whisper</center>")

demo.launch()