# --- START OF MODIFIED app.py --- try: import spaces print("'spaces' module imported successfully.") except ImportError: print("Warning: 'spaces' module not found. Using dummy decorator for local execution.") class DummySpaces: def GPU(self, *args, **kwargs): def decorator(func): print(f"Note: Dummy @GPU decorator used for function '{func.__name__}'.") return func return decorator spaces = DummySpaces() import gradio as gr from transformers import AutoTokenizer, AutoModelForSeq2SeqLM import torch import os import math # --- Configuration --- MODEL_PATH = "gregniuki/mandarin_thai_ipa" BATCH_SIZE = 8 # --- Device Setup --- if torch.cuda.is_available(): device = torch.device("cuda") print("GPU detected. Using CUDA.") else: device = torch.device("cpu") print("No GPU detected. Using CPU.") # --- Model Loading --- HF_AUTH_TOKEN = os.getenv("HF_TOKEN") print(f"Loading model and tokenizer from: {MODEL_PATH}") try: tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, token=HF_AUTH_TOKEN) model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_PATH, token=HF_AUTH_TOKEN) model.to(device) model.eval() print(f"Model and tokenizer loaded successfully on device: {device}") except Exception as e: raise RuntimeError(f"FATAL Error loading model/tokenizer: {e}") # --- Helper Function for Chunking --- def chunk_text(text, max_chars): if not text or text.isspace(): return [] text = text.strip() if len(text) <= max_chars: return [text] chunks, current_index = [], 0 while current_index < len(text): potential_end_index = min(current_index + max_chars, len(text)) actual_end_index = potential_end_index if potential_end_index < len(text): punctuation = ".!?。！？,，" best_split_pos = -1 for punc in punctuation: best_split_pos = max(best_split_pos, text.rfind(punc, current_index, potential_end_index)) if best_split_pos != -1: actual_end_index = best_split_pos + 1 chunk = text[current_index:actual_end_index] if chunk and not chunk.isspace(): chunks.append(chunk.strip()) current_index = actual_end_index if current_index >= len(text): break return [c for c in chunks if c] # --- Main Processing Function --- @spaces.GPU def translate_batch( text_input, max_chars_per_chunk, repetition_penalty, token_multiplier, # --- MODIFICATION: Added new parameter for the multiplier --- decoding_strategy, num_beams, length_penalty, use_early_stopping, temperature, top_p, progress=gr.Progress(track_tqdm=True) ): if not text_input or text_input.strip() == "": return "[Error] Please enter some text to process." lines = [line.strip() for line in text_input.splitlines() if line.strip()] if not lines: return "[Info] No valid text lines found in input." max_chars_per_chunk = int(max_chars_per_chunk) all_chunks = [] for line in lines: all_chunks.extend(chunk_text(line, max_chars_per_chunk)) if not all_chunks: return "[Info] No text chunks generated after processing input." generation_kwargs = {"repetition_penalty": repetition_penalty, "do_sample": False} if decoding_strategy == "Beam Search": generation_kwargs.update({ "num_beams": int(num_beams), "length_penalty": length_penalty, "early_stopping": use_early_stopping, }) elif decoding_strategy == "Sampling": generation_kwargs.update({ "do_sample": True, "temperature": temperature, "top_p": top_p, "num_beams": 1, }) print(f"Processing {len(all_chunks)} chunks with strategy: {decoding_strategy}. Args: {generation_kwargs}") print(f" Using token_multiplier: {token_multiplier}") # Log the multiplier all_ipa_outputs = [] num_batches = math.ceil(len(all_chunks) / BATCH_SIZE) for i in progress.tqdm(range(num_batches), desc="Processing Batches"): batch_start, batch_end = i * BATCH_SIZE, (i + 1) * BATCH_SIZE batch_chunks = all_chunks[batch_start:batch_end] try: inputs = tokenizer( batch_chunks, return_tensors="pt", padding=True, truncation=True, max_length=512 ).to(device) # --- MODIFICATION: Dynamic max_new_tokens using the slider value --- max_input_length = inputs["input_ids"].shape[1] generation_kwargs["max_new_tokens"] = min(int(max_input_length * token_multiplier) + 10, 512) with torch.no_grad(): outputs = model.generate(**generation_kwargs, **inputs) batch_ipa = tokenizer.batch_decode(outputs, skip_special_tokens=True) all_ipa_outputs.extend(batch_ipa) except Exception as e: print(f"Error during batch {i+1} processing: {e}") all_ipa_outputs.extend([f"[Error in batch {i+1}]"] * len(batch_chunks)) return "\n".join(all_ipa_outputs) # --- UI Helper Function for Dynamic Controls --- def update_decoding_ui(strategy): if strategy == "Beam Search": return gr.update(visible=True), gr.update(visible=False) elif strategy == "Sampling": return gr.update(visible=False), gr.update(visible=True) # --- Gradio UI using Blocks --- with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown( """ # 🇹🇭🇨🇳 Advanced Mandarin & Thai to IPA Converter Get the International Phonetic Alphabet (IPA) for Chinese (Mandarin) or Thai text. The model automatically detects the language. This interface provides advanced controls for tuning the output quality. """ ) with gr.Row(): with gr.Column(scale=1): input_textbox = gr.Textbox(lines=10, label="Input Text (Mandarin or Thai)", placeholder="Enter text here (e.g., 你好世界 or สวัสดีครับ).") submit_button = gr.Button("Submit", variant="primary") with gr.Column(scale=1): output_textbox = gr.Textbox(lines=10, label="IPA Output", interactive=False) with gr.Accordion("Generation & Chunking Controls", open=False): with gr.Row(): max_chars_slider = gr.Slider(minimum=1, maximum=512, step=1, value=36, label="Max Characters per Chunk", info="Splits long lines into smaller pieces for the model.") repetition_penalty_slider = gr.Slider(minimum=1.0, maximum=2.0, step=0.05, value=1.1, label="Repetition Penalty", info=">1.0 discourages repeated words. Prevents stuttering.") # --- MODIFICATION: Added the new slider for token multiplier --- token_multiplier_slider = gr.Slider( minimum=1.0, maximum=4.0, step=0.1, value=2.0, label="Output Token Multiplier", info="Safety factor for output length. Increase if output is cut off." ) decoding_strategy_radio = gr.Radio(["Beam Search", "Sampling"], value="Beam Search", label="Decoding Strategy", info="Choose the method for generating text.") with gr.Group(visible=True) as beam_search_group: with gr.Row(): num_beams_slider = gr.Slider(minimum=1, maximum=10, step=1, value=4, label="Number of Beams", info="More beams can yield better quality but are slower.") length_penalty_slider = gr.Slider(minimum=0.5, maximum=2.0, step=0.1, value=1.0, label="Length Penalty", info=">1.0 encourages longer output; <1.0 for shorter.") early_stopping_checkbox = gr.Checkbox(value=True, label="Early Stopping", info="Stop when a sentence is complete. Recommended.") with gr.Group(visible=False) as sampling_group: with gr.Row(): temperature_slider = gr.Slider(minimum=0.1, maximum=1.5, step=0.05, value=0.8, label="Temperature", info="Controls randomness. Lower is more predictable.") top_p_slider = gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.95, label="Top-p (Nucleus Sampling)", info="Considers a smaller, more probable set of next words.") gr.Markdown(f"**Model:** `{MODEL_PATH}` | **Batch Size:** `{BATCH_SIZE}` | **Device:** `{str(device).upper()}`") # --- Event Listeners --- decoding_strategy_radio.change(fn=update_decoding_ui, inputs=decoding_strategy_radio, outputs=[beam_search_group, sampling_group]) submit_button.click( fn=translate_batch, inputs=[ input_textbox, max_chars_slider, repetition_penalty_slider, token_multiplier_slider, # --- MODIFICATION: Added slider to inputs list --- decoding_strategy_radio, num_beams_slider, length_penalty_slider, early_stopping_checkbox, temperature_slider, top_p_slider, ], outputs=output_textbox ) # --- Launch the App --- if __name__ == "__main__": demo.launch() # --- END OF FILE ---