Update app.py

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
 import torch
 import pickle
-from model import predict_with_beam_search, PosFormerImprovedWithIAC, PAD_TOKEN, SOS_TOKEN, EOS_TOKEN
+from model import predict_with_beam_search, PosFormerImprovedWithIAC, PAD_TOKEN, SOS_TOKEN, EOS_TOKEN, UNK_TOKEN
 
-# --- Загрузка модели и vocab ---
+# --- загрузка модели (как в предыдущем примере) ---
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 with open("vocab.pkl", "rb") as f:
     vocab = pickle.load(f)
@@ -16,24 +16,28 @@ model = PosFormerImprovedWithIAC(
     id_to_token_map={i:s for s,i in vocab.stoi.items()}
 ).to(DEVICE)
 
-checkpoint = torch.load("best_model.pth", map_location=DEVICE)
+checkpoint = torch.load("d256 layers=4 densenet-121 32%.pth", map_location=DEVICE)
 model.load_state_dict(checkpoint["model_state_dict"])
 model.eval()
 
-# --- Функция для Gradio ---
+# --- функция инференса ---
 def inference(image, beam_size=5):
+    # image здесь — PIL.Image из Sketchpad
     return predict_with_beam_search(model, image, vocab, beam_size, DEVICE)
 
-# --- Gradio Interface ---
+# --- интерфейс с полем рисования ---
 iface = gr.Interface(
     fn=inference,
     inputs=[
-        gr.inputs.Image(type="pil", label="Upload Image"),
-        gr.inputs.Slider(minimum=1, maximum=20, default=10, label="Beam Size")
+        gr.Sketchpad(
+            shape=(IMG_HEIGHT, IMG_WIDTH), 
+            label="Нарисуйте формулу"
+        ),
+        gr.Slider(minimum=1, maximum=20, default=10, label="Beam Size")
     ],
-    outputs=gr.outputs.Textbox(label="Predicted LaTeX"),
+    outputs=gr.Textbox(label="Предсказанная LaTeX-строка"),
     title="Handwritten Formula Recognition",
-    description="Загрузите изображение, и модель вернёт LaTeX-строку."
+    description="Нарисуйте формулу на холсте и получите LaTeX."
 )
 
 if __name__ == "__main__":