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 = """
/* Import a nice font (Poppins) */
@import url('https://fonts.googleapis.com/css2?family=Poppins:wght@300;400;600&display=swap');

/* Apply font to everything */
body, button, input, select, textarea, .gradio-container {
    font-family: 'Poppins', sans-serif !important;
}

/* Header Styling */
.header-title {
    text-align: center;
    color: #2c3e50;
    margin-bottom: 0.5rem;
    font-size: 2.2em;
    font-weight: 600;
}
.header-subtitle {
    text-align: center;
    color: #7f8c8d;
    margin-bottom: 2rem;
    font-weight: 300;
    font-size: 1.1em;
}

/* Card Container Styling */
.input-container, .output-container {
    background: #ffffff;
    border: 1px solid #e0e0e0;
    border-radius: 12px;
    padding: 25px;
    box-shadow: 0 4px 6px rgba(0,0,0,0.05);
    margin-bottom: 20px;
}

/* Button Styling */
#translate-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    border: none;
    color: white;
    font-weight: 600;
    font-size: 1.1em;
    padding: 10px 20px;
    border-radius: 8px;
    transition: transform 0.1s;
}
#translate-btn:hover {
    transform: scale(1.02);
    box-shadow: 0 5px 15px rgba(0,0,0,0.2);
}

/* Feedback Section */
.feedback-box {
    background-color: #f8f9fa;
    border-radius: 8px;
    padding: 20px;
    border-left: 5px solid #ffc107;
}
"""

# --- THE UI ---
with gr.Blocks(css=custom_css, title="Kurukh AI Translator") as demo:
    
    # HEADER SECTION
    with gr.Row():
        with gr.Column():
            gr.Markdown("# 🇮🇳 AI Kurukh (Oraon) Translator", elem_classes=["header-title"])
            gr.Markdown("### Bridging Communities with Artificial Intelligence | Voice & Hinglish Supported", elem_classes=["header-subtitle"])

    # TABS FOR NAVIGATION
    with gr.Tabs():
        
        # --- TAB 1: TRANSLATOR INTERFACE ---
        with gr.TabItem("🗣️ Translator"):
            
            # Help / Instructions (Collapsible)
            with gr.Accordion("ℹ️ How to use (Click to expand)", open=False):
                gr.Markdown("""
                1. **Select Mode:** Choose translation direction (Kurukh -> Hindi OR Hindi -> Kurukh).
                2. **Hinglish Support:** Check the box if you are typing Hindi words using English letters (e.g., 'Tumhara naam kya hai').
                3. **Input Methods:** Type text in the box OR click the microphone to speak (Hindi audio supported).
                4. **Results:** View the translation and listen to the audio output (for Hindi results).
                """)

            # Main Grid Layout
            with gr.Row():
                # LEFT COLUMN (INPUT CARD)
                with gr.Column(elem_classes=["input-container"]):
                    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,
                        show_label=False
                    )
                    
                    input_audio = gr.Audio(
                        sources=["microphone"], 
                        type="filepath", 
                        label="🎙️ Voice Input (Hindi Only)"
                    )
                    
                    translate_btn = gr.Button("Translate 🚀", variant="primary", elem_id="translate-btn")

                # RIGHT COLUMN (OUTPUT CARD)
                with gr.Column(elem_classes=["output-container"]):
                    gr.Markdown("### 📤 Translation Result")
                    
                    output_text = gr.Textbox(
                        label="Translation", 
                        lines=4, 
                        show_label=False,
                        interactive=False,
                        show_copy_button=True
                    )
                    
                    output_audio = gr.Audio(
                        label="🔊 Listen (Hindi Only)", 
                        interactive=False,
                        autoplay=False
                    )

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

        # --- TAB 2: FEEDBACK INTERFACE ---
        with gr.TabItem("📝 Improve the AI"):
            with gr.Column(elem_classes=["feedback-box"]):
                gr.Markdown("### 🛠️ Help us improve accuracy")
                gr.Markdown("Our AI is still learning! If you spot a mistake, please submit the correct translation below. This helps us train better versions.")
                
                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", placeholder="The sentence you tried to translate")
                    fb_ai_output = gr.Textbox(label="AI's Wrong Translation (Optional)")
                
                fb_user_correct = gr.Textbox(label="✅ Your Correct Translation", lines=2, placeholder="Type the correct Kurukh/Hindi 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
                )
    
    # FOOTER
    gr.Markdown("---")
    gr.Markdown("<center>Built with ❤️ for the Kurukh Community • Powered by Google mT5 & OpenAI Whisper</center>")

demo.launch()