Update app.py

app.py

@@ -1,9 +1,8 @@
-import torch
 import gradio as gr
 from transformers import pipeline, AutoTokenizer
 from functools import lru_cache
 
-MODEL_ID = "tasal9/ZamAI-mT5-Pashto"  # Must be a string
+MODEL_ID = "tasal9/ZamAI-mT5-Pashto"
 
 SAMPLE_INSTRUCTIONS = [
     "په پښتو کې د خپل نوم او د عمر معلومات ولیکئ.",
@@ -17,53 +16,52 @@ SAMPLE_INSTRUCTIONS = [
 
 def generate_prompt(instruction, input_text=""):
     if input_text:
-        return f"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:"
+        return f"### Instruction:\n{instruction}\n\n### Input:\n{input_text}\n\n### Response:"
     else:
-        return f"Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n### Instruction:\n{instruction}\n\n### Response:"
+        return f"### Instruction:\n{instruction}\n\n### Response:"
 
-# Cache the pipeline for performance
 @lru_cache(maxsize=1)
 def get_generator():
-    model_id_str = str(MODEL_ID)  # Ensure it's a string
-    # Optional: verify tokenizer exists
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(model_id_str, use_fast=False)
-    except Exception as e:
-        raise ValueError(f"Failed to load tokenizer for {model_id_str}: {e}")
-
-    device = 0 if torch.cuda.is_available() else -1
+    # Always CPU in ZeroGPU
+    AutoTokenizer.from_pretrained(MODEL_ID, use_fast=False)
     return pipeline(
         "text2text-generation",
-        model=model_id_str,
-        tokenizer=model_id_str,
-        device=device,
+        model=MODEL_ID,
+        tokenizer=MODEL_ID,
+        device=-1,                 # Force CPU
+        return_full_text=False,    # Avoids prompt repetition
         use_fast=False
     )
 
 def predict(instruction, input_text, max_length, num_beams, temperature, top_p):
     gen = get_generator()
     prompt = generate_prompt(instruction, input_text)
+
     outputs = gen(
         prompt,
         max_length=max_length,
         num_beams=num_beams,
         temperature=temperature,
         top_p=top_p,
+        do_sample=True,       # Sampling works better on CPU than beams
         early_stopping=True
     )
-    # Remove splitting on 'ځواب:' and just return the generated text after the prompt
+
     generated = outputs[0]["generated_text"].strip()
-    # Remove the prompt from the generated text if present
-    if generated.startswith(prompt):
-        generated = generated[len(prompt):].strip()
-    return generated if generated else "No response generated."
 
-# Build UI with Blocks
+    # Clean output
+    for cut in [prompt, "### Instruction:", "### Response:", "ځواب:"]:
+        if generated.startswith(cut):
+            generated = generated[len(cut):].strip()
+
+    return generated if generated else "⚠️ No response generated."
+
+# ---------------- Gradio UI ----------------
 with gr.Blocks() as demo:
     gr.Markdown(
         """
         # ZamAI mT5 Pashto Demo
-         اپلیکیشن  **ZamAI-mT5-Pashto** ماډل د پښتو لارښوونو لپاره .
+        اپلیکیشن  **ZamAI-mT5-Pashto** د پښتو لارښوونو لپاره.  
         لاندې تنظیمات بدل کړئ او لارښوونه ولیکئ ترڅو ځواب ترلاسه کړئ.
         """
     )
@@ -72,7 +70,7 @@ with gr.Blocks() as demo:
         with gr.Column(scale=2):
             instruction_dropdown = gr.Dropdown(
                 choices=SAMPLE_INSTRUCTIONS,
-                label="نمونې لارښوونې (Auto-select)",
+                label="نمونې لارښوونې",
                 value=SAMPLE_INSTRUCTIONS[0],
                 interactive=True
             )
@@ -84,14 +82,17 @@ with gr.Blocks() as demo:
             input_text = gr.Textbox(lines=2, placeholder="اختیاري متن...", label="متن")
             output = gr.Textbox(label="ځواب", interactive=False, lines=5)
             generate_btn = gr.Button("جوړول", variant="primary")
+
         with gr.Column(scale=1):
-            gr.Markdown("### د تولید تنظیمات")
-            max_length = gr.Slider(32, 512, value=256, step=1, label="اعظمي اوږدوالی")
-            num_beams = gr.Slider(1, 10, value=5, step=1, label="شمیر شعاعونه")
-            temperature = gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="تودوخه")
-            top_p = gr.Slider(0.1, 1.0, value=1.0, step=0.05, label="Top-p")
+            gr.Markdown("### د تولید تنظیمات (ZeroGPU)")
+            max_length = gr.Slider(32, 256, value=128, step=1, label="اعظمي اوږدوالی")
+            num_beams = gr.Slider(1, 3, value=2, step=1, label="شمیر شعاعونه")
+            temperature = gr.Slider(0.5, 1.5, value=1.0, step=0.1, label="تودوخه")
+            top_p = gr.Slider(0.7, 1.0, value=0.9, step=0.05, label="Top-p")
 
-    instruction_dropdown.change(lambda x: x, inputs=instruction_dropdown, outputs=instruction_textbox)
+    instruction_dropdown.change(
+        lambda x: x, inputs=instruction_dropdown, outputs=instruction_textbox
+    )
 
     generate_btn.click(
         fn=predict,
@@ -100,4 +101,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()