Update Main.py

Main.py

@@ -5,10 +5,9 @@ import torch
 from transformers import LightOnOcrForConditionalGeneration, LightOnOcrProcessor
 from PIL import Image
 
-# Ускоряем скачивание модели
+# Ускоряем скачивание на HF Spaces
 os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 
-# ==================== Настройки страницы ====================
 st.set_page_config(
     page_title="LightOnOCR • Распознай текст",
     page_icon="📄",
@@ -16,30 +15,22 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
-# ==================== Кастомный CSS ====================
+# Простой CSS
 st.markdown("""
-    <style>
-        .main { 
-            background: linear-gradient(180deg, #f8f9fa, #e9f0f7); 
-        }
-        .header-emoji { 
-            font-size: 3.5rem; 
-            text-align: center; 
-            margin: 15px 0; 
-        }
-        .result-box {
-            background: #ffffff;
-            border-radius: 16px;
-            padding: 24px;
-            box-shadow: 0 10px 30px rgba(0, 0, 0, 0.08);
-            border: 1px solid #e5e7eb;
-            margin-top: 20px;
-        }
-    </style>
+<style>
+    .main { background: linear-gradient(180deg, #f8f9fa, #e9f0f7); }
+    .result-box {
+        background: #ffffff;
+        border-radius: 16px;
+        padding: 24px;
+        box-shadow: 0 10px 30px rgba(0,0,0,0.08);
+        margin-top: 20px;
+    }
+    .header-emoji { font-size: 3.5rem; text-align: center; margin: 15px 0; }
+</style>
 """, unsafe_allow_html=True)
 
-# ==================== Загрузка модели ====================
-@st.cache_resource(show_spinner="⏳ Загрузка модели LightOnOCR-1B-1025...\nЭто может занять 2–6 минут при первом запуске...")
+@st.cache_resource(show_spinner="⏳ Загрузка модели LightOnOCR-1B-1025...\nЭто может занять 2–6 минут при первом запуске на CPU")
 def load_model():
     model_name = "lightonai/LightOnOCR-1B-1025"
     
@@ -54,13 +45,20 @@ def load_model():
 
     processor = LightOnOcrProcessor.from_pretrained(model_name)
     
-    # Исправляем pad_token
-    if processor.tokenizer.pad_token is None:
-        processor.tokenizer.pad_token = processor.tokenizer.eos_token
-
     return processor, model, device, dtype
 
-# ==================== Загрузка изображения ====================
+# ====================== Заголовок ======================
+st.markdown('<div class="header-emoji">📄✨</div>', unsafe_allow_html=True)
+st.title("LightOnOCR")
+st.markdown("**Распознавание текста с изображений**")
+st.caption("Модель: lightonai/LightOnOCR-1B-1025")
+
+# ====================== Загрузка модели ======================
+processor, model, device, dtype = load_model()
+
+st.sidebar.success(f"✅ Модель загружена на **{device.upper()}**")
+
+# ====================== Загрузка изображения ======================
 def load_image():
     uploaded_file = st.file_uploader(
         "📸 Загрузите изображение (png, jpg, jpeg, webp)",
@@ -72,71 +70,55 @@ def load_image():
         return Image.open(io.BytesIO(image_data)).convert('RGB')
     return None
 
-# ==================== Основной интерфейс ====================
-st.markdown('<div class="header-emoji">📄✨</div>', unsafe_allow_html=True)
-st.title("LightOnOCR")
-st.markdown("**Мгновенное распознавание текста с изображений**")
-st.caption("Модель: lightonai/LightOnOCR-1B-1025")
-
-# Загрузка модели
-processor, model, device, dtype = load_model()
-
-# Сайдбар
-with st.sidebar:
-    st.success(f"✅ Модель загружена на **{device.upper()}**")
-    st.info("Лучше всего работает с английским текстом, документами и таблицами.")
-
-# Загрузка изображения
 img = load_image()
 
-# ==================== Распознавание ====================
+# ====================== Распознавание ======================
 if st.button("🔍 Распознать текст", use_container_width=True, type="primary"):
     if img is None:
-        st.error("Пожалуйста, сначала загрузите изображение")
+        st.error("Сначала загрузите изображение")
     else:
-        with st.spinner("Распознавание текста... (может занять 5–20 секунд на CPU)"):
+        with st.spinner("Распознавание текста... (может занять 5–20 сек на CPU)"):
             
-            # Правильный промпт
+            # ✅ Правильный формат для LightOnOCR (только изображение + промпт)
             conversation = [
                 {
                     "role": "user",
                     "content": [
-                        {"type": "image"},
-                        {"type": "text", "text": "Extract all the text from this image accurately. Preserve original formatting, tables, and line breaks as much as possible."}
+                        {"type": "image"},   # изображение передаётся автоматически через processor
+                        {"type": "text", "text": "Extract all the text from this image as accurately as possible. Output clean text with preserved line breaks and formatting."}
                     ]
                 }
             ]
 
-            # Подготовка входных данных
+            # Подготовка inputs (processor обработает и изображение, и текст)
             inputs = processor.apply_chat_template(
                 conversation,
                 add_generation_prompt=True,
                 tokenize=True,
                 return_dict=True,
                 return_tensors="pt",
+                # Важно: передаём само PIL-изображение
                 images=img
             )
 
-            # Перенос на устройство
-
-
-inputs = {
+            # Переносим на устройство
+            inputs = {
                 k: (v.to(device=device, dtype=dtype) if v.is_floating_point() else v.to(device))
                 for k, v in inputs.items()
             }
 
             # Генерация
-        output_ids = model.generate(
+            output_ids = model.generate(
                 **inputs,
                 max_new_tokens=2048,
                 do_sample=False,
                 temperature=0.0,
-                num_beams=1,
+                num_beams=1,          # для стабильности
                 pad_token_id=processor.tokenizer.pad_token_id,
                 eos_token_id=processor.tokenizer.eos_token_id,
             )
 
-            # Убираем промпт из результата
+            # Убираем входной промпт — оставляем только сгенерированный текст
             prompt_length = inputs["input_ids"].shape[1]
             generated_ids = output_ids[0, prompt_length:]
             
@@ -146,20 +128,19 @@ inputs = {
                 clean_up_tokenization_spaces=True
             ).strip()
 
-            # Результат
+            # Вывод результата
             st.success("✅ Распознавание завершено!")
             st.markdown('<div class="result-box">', unsafe_allow_html=True)
             st.subheader("📝 Распознанный текст")
-            st.code(generated_text, language=None)
+            st.code(generated_text, language=None)   # лучше чем markdown для больших блоков
             st.markdown('</div>', unsafe_allow_html=True)
 
             st.download_button(
-                label="💾 Скачать текст как .txt",
+                label="💾 Скачать как .txt",
                 data=generated_text,
                 file_name="recognized_text.txt",
                 mime="text/plain"
             )
 
-# Подвал
 st.markdown("---")
 st.caption("Сделано на базе [lightonai/LightOnOCR-1B-1025](https://huggingface.co/lightonai/LightOnOCR-1B-1025)")