Update app.py

app.py

@@ -1,35 +1,80 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
 
-MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # در صورت نیاز مدل کوچک‌تر انتخاب کنید
 
+# بارگذاری مدل و توکنایزر
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    device_map="auto",           # GPU اگر باشد استفاده می‌کند
+    torch_dtype="auto"           # دقت مناسب سخت‌افزار
+)
 
+# ساخت پایپلاین تولید متن
 generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 def generate_text(prompt, max_new_tokens, temperature, top_p):
-    result = generator(
-        prompt,
-        max_new_tokens=int(max_new_tokens),
-        do_sample=True,
-        temperature=float(temperature),
-        top_p=float(top_p),
-    )[0]["generated_text"]
-    return result
+    """
+    تولید پاسخ با استفاده از max_new_tokens (مستقل از طول ورودی).
+    در صورت خطای حافظه، یک بار با نصف max_new_tokens تلاش مجدد می‌کند.
+    """
+    max_new_tokens = int(max_new_tokens)
+    temperature = float(temperature)
+    top_p = float(top_p)
+
+    def _run(tokens):
+        return generator(
+            prompt,
+            max_new_tokens=tokens,  # مستقل از طول ورودی
+            do_sample=True,
+            temperature=temperature,
+            top_p=top_p,
+        )[0]["generated_text"]
+
+    try:
+        return _run(max_new_tokens)
+    except torch.cuda.OutOfMemoryError:
+        # تلاش مجدد با نصف مقدار
+        try:
+            torch.cuda.empty_cache()
+            smaller = max(32, max_new_tokens // 2)
+            return _run(smaller) + "\n\n[⚠️ حافظه کم بود؛ خروجی با توکن‌های کمتر تولید شد. اگر باز هم مشکل داری، مقدار max_new_tokens را کمتر کن یا از مدل کوچک‌تر استفاده کن.]"
+        except Exception:
+            torch.cuda.empty_cache()
+            return "⚠️ خطای حافظه (OOM). لطفاً مقدار max_new_tokens را کاهش بده یا از مدل کوچک‌تر/سخت‌افزار قوی‌تر استفاده کن."
+    except RuntimeError as e:
+        # برخی OOMها RuntimeError هستند
+        if "out of memory" in str(e).lower():
+            try:
+                torch.cuda.empty_cache()
+                smaller = max(32, max_new_tokens // 2)
+                return _run(smaller) + "\n\n[⚠️ حافظه کم بود؛ خروجی با توکن‌های کمتر تولید شد. اگر باز هم مشکل داری، مقدار max_new_tokens را کمتر کن یا از مدل کوچک‌تر استفاده کن.]"
+            except Exception:
+                torch.cuda.empty_cache()
+                return "⚠️ خطای حافظه (OOM). لطفاً مقدار max_new_tokens را کاهش بده یا از مدل کوچک‌تر/سخت‌افزار قوی‌تر استفاده کن."
+        # سایر خطاها
+        return f"❌ خطا: {e}"
+    except Exception as e:
+        return f"❌ خطا: {e}"
 
 with gr.Blocks(title="TinyLlama Chat") as demo:
-    gr.Markdown("## TinyLlama Chat\nگوش به زنگ م.")
+    gr.Markdown(
+        "## TinyLlama Chat\n"
+        "پرامپتت را بنویس و پاسخ بگیر.\n\n"
+        "**نکته:** اگر خطای حافظه دیدی، مقدار `max_new_tokens` را کم کن یا مدل کوچک‌تری انتخاب کن."
+    )
     with gr.Row():
-        prompt = gr.Textbox(label="Prompt", lines=4, placeholder="مثال: درباره گل‌های آفریقا بگو...")
+        prompt = gr.Textbox(label="Prompt", lines=4, placeholder="مثال: درباره گل‌های آفریقا بگو…")
     with gr.Row():
         max_new_tokens = gr.Slider(16, 512, value=150, step=1, label="max_new_tokens")
         temperature = gr.Slider(0.1, 1.5, value=0.7, step=0.05, label="temperature")
         top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p")
-    btn = gr.Button("جواب بده")
-    output = gr.Textbox(label="یاور : ", lines=10)
+    btn = gr.Button("تولید پاسخ")
+    output = gr.Textbox(label="خروجی", lines=12)
 
     btn.click(generate_text, [prompt, max_new_tokens, temperature, top_p], output)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()