updated app.py for greedy

app.py

@@ -4,6 +4,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 MODEL_ID = "Velkamez/tamazight-sp-bpe"  # or "./model"
 
+# ---- Load tokenizer & model ----
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID,
     trust_remote_code=True
@@ -11,38 +12,43 @@ tokenizer = AutoTokenizer.from_pretrained(
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float32,
     trust_remote_code=True
 )
 model.eval()
 
 
+# ---- Generation function ----
 def generate_reply(message):
+    # Encode input
     inputs = tokenizer(message, return_tensors="pt")
 
     with torch.no_grad():
         output_ids = model.generate(
             **inputs,
             max_new_tokens=64,
-            do_sample=False,
+            do_sample=False,  # greedy
             pad_token_id=tokenizer.eos_token_id
         )
 
-    decoded = tokenizer.decode(
-        output_ids[0],
-        skip_special_tokens=True
-    )
+    # 🔑 remove prompt tokens
+    input_len = inputs["input_ids"].shape[1]
+    new_token_ids = output_ids[0][input_len:]
 
-    if decoded.startswith(message):
-        decoded = decoded[len(message):].strip()
+    # 🔑 tokenizer-aware decoding (fixes "az ul")
+    tokens = tokenizer.convert_ids_to_tokens(new_token_ids)
+    decoded = tokenizer.convert_tokens_to_string(tokens)
 
-    return decoded
+    return decoded.strip()
 
 
+# ---- UI ----
 with gr.Blocks() as demo:
     gr.Markdown("# 🗣️ Tamazight LLM Demo")
 
-    msg = gr.Textbox(label="Message")
+    msg = gr.Textbox(
+        label="Message",
+        placeholder="Write something (e.g. azul)"
+    )
     out = gr.Textbox(label="Model output")
 
     msg.submit(generate_reply, msg, out)