Update app.py

app.py

@@ -1,9 +1,8 @@
 import gradio as gr
 import easyocr
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 from PIL import Image
 import numpy as np
-import os
 from typing import Tuple
 
 ## 1. تنظیمات اولیه و مدل‌ها
@@ -29,17 +28,21 @@ class TextPostProcessor:
             '۵':'5', '۶':'6', '۷':'7', '۸':'8', '۹':'9'
         }
         
-        # بارگذاری مدل زبانی
+        # بارگذاری مدل زبانی فارسی HooshvareLab/gpt2-fa
         try:
-            self.llm = pipeline("text-generation", model="gpt2")
-        except:
+            self.llm = pipeline(
+                "text-generation", 
+                model="HooshvareLab/gpt2-fa", 
+                tokenizer="HooshvareLab/gpt2-fa"
+            )
+        except Exception as e:
             self.llm = None
+            print(f"خطا در بارگذاری مدل HooshvareLab/gpt2-fa: {str(e)}")
     
     def preprocess(self, text: str) -> str:
         """پیش‌پردازش متن استخراج شده"""
         if not text:
             return ""
-            
         # نرمال‌سازی متن
         for old, new in self.replacements.items():
             text = text.replace(old, new)
@@ -49,15 +52,21 @@ class TextPostProcessor:
         """بهبود متن با مدل زبانی"""
         if not text or not self.llm:
             return text
-            
         try:
+            prompt = f"متن زیر را ویرایش و روان‌تر کن:\n{text}\n"
             enhanced = self.llm(
-                f"اصلاح و بازنویسی متن فارسی زیر:\n{text}\n\nمتن بهبود یافته:",
-                max_length=200,
-                num_return_sequences=1
+                prompt,
+                max_length=len(prompt) + 60,
+                num_return_sequences=1,
+                do_sample=True,
+                temperature=0.8,
+                pad_token_id=0
             )
-            return enhanced[0]['generated_text']
-        except:
+            # فقط متن تولیدی پس از پرامپت را برگردان
+            gen_text = enhanced[0]['generated_text']
+            return gen_text[len(prompt):].strip() if gen_text.startswith(prompt) else gen_text.strip()
+        except Exception as e:
+            print(f"خطا در بهبود متن با LLM: {str(e)}")
             return text
 
 ## 2. پردازش اصلی