Spaces:
Sleeping
Sleeping
| import spaces # MUST be the very first import for ZeroGPU Spaces | |
| import os | |
| import gradio as gr | |
| import torch | |
| import regex | |
| import unicodedata | |
| import re | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| from peft import PeftModel | |
| # ─── Model Configuration ────────────────────────────────────────────────────── | |
| BASE_MODEL = "unsloth/gemma-3-270m-it" | |
| LORA_PATH = "Razon2006/gemma3-tamil-Grapheme-NewChange-r32-lr" | |
| hf_token = os.environ.get("HF_TOKEN") | |
| print("Loading tokenizer globally...") | |
| tokenizer = AutoTokenizer.from_pretrained(LORA_PATH, token=hf_token) | |
| # Global model variable for lazy loading | |
| model = None | |
| # ─── Grapheme Tokenizer & Constants ─────────────────────────────────────────── | |
| INSTRUCTION = ( | |
| "கீழ்க்காணும் தமிழ் வெண்பாவை, அதன் சொற்களை சரியான எல்லைகளுடன் " | |
| "பிரித்து எழுதுக (பதப்பிரிவு). சொற்களின் வரிசை மாறாமல், " | |
| "ஒட்டிய சொற்களை பிரிக்கவும், தேவையெனில் பிரிந்த எழுத்துக்களை " | |
| "இணைத்து சரியான சொல்லை உருவாக்கவும். பதப்பிரிவை மட்டும் தரவும்." | |
| ) | |
| PROBLEMATIC_AKSHARAS = { | |
| 'யு', 'றா', 'ஞா', 'யே', 'யை', 'றே', 'கெ', 'றோ', 'யெ', 'யொ', | |
| 'றெ', 'ழா', 'நூ', 'லெ', 'வொ', 'ரெ', 'றொ', 'கீ', 'டே', 'லொ', | |
| 'னொ', 'டொ', 'ணே', 'ளெ', 'ணோ', 'ரொ', 'றூ', 'ணீ', 'றீ', 'ழீ', | |
| 'னூ', 'ணெ', 'வூ', 'யீ', 'டீ', 'ளொ', 'லூ', 'நொ', 'ணொ', 'டூ', | |
| 'நை', 'கௌ', 'ளூ', 'ழே', 'ழூ', 'ஞை', 'ளீ', 'ஜை', 'ணூ', 'ழொ', | |
| 'ழெ', 'ஞு', 'ஞெ', 'ஞூ', 'ஸா', 'ஜூ', 'எா', 'எீ', 'ழோ', 'பௌ', | |
| } | |
| def clean_grapheme_split(text, tokenizer): | |
| words = text.split(' ') | |
| result_words = [] | |
| for word in words: | |
| if not word.strip(): | |
| continue | |
| clusters = regex.findall(r'\X', word) | |
| clusters = [ | |
| unicodedata.normalize('NFC', c) for c in clusters | |
| if c.strip() and any('\u0B80' <= ch <= '\u0BFF' for ch in c) | |
| ] | |
| if not clusters: | |
| result_words.append(word) | |
| continue | |
| merged = [] | |
| i = 0 | |
| while i < len(clusters): | |
| c = clusters[i] | |
| tokens = tokenizer.tokenize(c) | |
| if (c in PROBLEMATIC_AKSHARAS or len(tokens) > 1) and i + 1 < len(clusters): | |
| combined = c + clusters[i+1] | |
| while len(tokenizer.tokenize(combined)) > 1 and i + 2 < len(clusters): | |
| i += 1 | |
| combined = combined + clusters[i+1] | |
| merged.append(combined) | |
| i += 2 | |
| else: | |
| merged.append(c) | |
| i += 1 | |
| result_words.append('|'.join(merged)) | |
| return ' '.join(result_words) | |
| def post_process_grapheme(text): | |
| return text.replace('|', '') | |
| # ─── Inference Pipeline ─────────────────────────────────────────────────────── | |
| def generate_pathaseth(venba, max_new_tokens=256): | |
| global model | |
| # LAZY LOADING: Load the model here where CUDA is fully unlocked by ZeroGPU | |
| if model is None: | |
| print("Loading base model directly into GPU...") | |
| base_model = AutoModelForCausalLM.from_pretrained( | |
| BASE_MODEL, | |
| torch_dtype=torch.bfloat16, | |
| device_map="cuda", # <--- FIX: Forces native loading to GPU to avoid precision loss | |
| token=hf_token | |
| ) | |
| print("Loading LoRA adapter...") | |
| model = PeftModel.from_pretrained(base_model, LORA_PATH, token=hf_token) | |
| model.eval() | |
| print("Model ready for inference!") | |
| venba = venba.strip() | |
| if not venba: | |
| return "Please enter a Venba verse." | |
| venba_clean = re.sub(r'\s+', ' ', venba).replace('்்', '்').strip() | |
| venba_no_space = venba_clean.replace(' ', '') | |
| venba_g = clean_grapheme_split(venba_no_space, tokenizer) | |
| messages = [{ | |
| "role": "user", | |
| "content": f"{INSTRUCTION}\n\nவெண்பா: {venba_g}\nபதப்பிரிவு:" | |
| }] | |
| prompt = tokenizer.apply_chat_template( | |
| messages, tokenize=False, add_generation_prompt=True | |
| ) | |
| bos = tokenizer.bos_token | |
| if bos and prompt.startswith(bos): | |
| prompt = prompt[len(bos):] | |
| inputs = tokenizer(prompt, return_tensors="pt").to("cuda") | |
| with torch.no_grad(): | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=max_new_tokens, | |
| use_cache=True, | |
| do_sample=False, # <--- FIX: Ensures deterministic output (stops hallucination) | |
| ) | |
| new_tokens = outputs[0][inputs['input_ids'].shape[1]:] | |
| raw = tokenizer.decode(new_tokens, skip_special_tokens=True).strip() | |
| return post_process_grapheme(raw) | |
| # ─── Example Venbas ─────────────────────────────────────────────────────────── | |
| EXAMPLES = [ | |
| ["இறுதி பயப்பினு மெஞ்சா திறைவற் குறுதி பயப்பதாந் தூது"], | |
| ] | |
| # ─── Gradio Interface ───────────────────────────────────────────────────────── | |
| with gr.Blocks( | |
| title="Tamil Venba Pathasetham", | |
| theme=gr.themes.Soft() | |
| ) as demo: | |
| gr.HTML(""" | |
| <div style="text-align: center; font-size: 2em; font-weight: bold; color: #4B0082;">Tamil Venba Pathasetham (பதசேதம்)</div> | |
| <div style="text-align: center; color: #666; margin-bottom: 20px;"> | |
| Automated word segmentation of classical Tamil Venba poetry using Gemma 3 + Grapheme Tokenization | |
| </div> | |
| """) | |
| with gr.Row(): | |
| with gr.Column(): | |
| venba_input = gr.Textbox( | |
| label="Enter Tamil Venba (வெண்பா)", | |
| placeholder="வண்மைதரு மாகமநூன் வைத்த பொருள்வழுவா...", | |
| lines=4, | |
| ) | |
| submit_btn = gr.Button("Generate Pathasetham →", variant="primary", size="lg") | |
| with gr.Column(): | |
| pathaseth_output = gr.Textbox( | |
| label="Pathasetham (பதசேதம்) Output", | |
| lines=4, | |
| interactive=False, | |
| ) | |
| gr.Examples( | |
| examples=EXAMPLES, | |
| inputs=venba_input, | |
| outputs=pathaseth_output, | |
| fn=generate_pathaseth, | |
| cache_examples=False, | |
| label="Click any example to test instantly" | |
| ) | |
| submit_btn.click( | |
| fn=generate_pathaseth, | |
| inputs=venba_input, | |
| outputs=pathaseth_output, | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch() |