Update app.py

app.py

@@ -1,26 +1,22 @@
 # app.py
-
 import gradio as gr
 import torch
 import numpy as np
 from transformers import AutoTokenizer, AutoModelForTokenClassification
 
 # ============================
-# 1. Модель мен токенизаторды жүктеу
+# 1. Модель мен токенизатор
 # ============================
-# Мысалы, мультилингвалды BERT (қазақ тілін қолдайды)
 model_checkpoint = "bert-base-multilingual-cased"
 tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
 model = AutoModelForTokenClassification.from_pretrained(model_checkpoint, num_labels=7)
 
-# Қазақ NER үшін label тізімі
 label_list = ["O", "B-PER", "I-PER", "B-ORG", "I-ORG", "B-LOC", "I-LOC"]
 
 # ============================
 # 2. NER функциясы
 # ============================
 def predict_ner(text):
-    # Токенизация
     tokens = tokenizer(text.split(), return_tensors="pt", is_split_into_words=True)
     outputs = model(**tokens).logits
     predictions = np.argmax(outputs.detach().numpy(), axis=2)[0]
@@ -33,24 +29,42 @@ def predict_ner(text):
             label = label_list[predictions[idx]]
             word = text.split()[word_idx]
             if label != "O":
-                results.append(f"{word} → {label}")
+                results.append((word, label))
             already_seen.add(word_idx)
-    if not results:
+    return results
+
+# ============================
+# 3. Түсті HTML шығару
+# ============================
+def highlight_ner(text):
+    entities = predict_ner(text)
+    if not entities:
         return "Атаулар табылған жоқ"
-    return "\n".join(results)
+
+    colored_text = text
+    for word, label in entities:
+        color = ""
+        if label.startswith("PER"):
+            color = "red"
+        elif label.startswith("ORG"):
+            color = "blue"
+        elif label.startswith("LOC"):
+            color = "green"
+        colored_text = colored_text.replace(word, f"<span style='color:{color}; font-weight:bold'>{word}</span>")
+    return colored_text
 
 # ============================
-# 3. Gradio интерфейсі
+# 4. Gradio интерфейсі
 # ============================
 iface = gr.Interface(
-    fn=predict_ner,
-    inputs=gr.Textbox(lines=5, placeholder="Мәтінді осында енгізіңіз..."),
-    outputs=gr.Textbox(label="Анықталған атаулар (NER)"),
+    fn=highlight_ner,
+    inputs=gr.Textbox(lines=10, placeholder="Қазақ мәтінін осында енгізіңіз..."),
+    outputs=gr.HTML(label="Анықталған атаулар"),
     title="Қазақ тіліндегі NER",
-    description="Бұл құрал қазақ мәтіндеріндегі адам, ұйым және орын атауларын анықтайды"
+    description="PER – адам, ORG – ұйым, LOC – орын. Атаулар мәтінде түспен ерекшеленеді."
 )
 
 # ============================
-# 4. Іске қосу
+# 5. Іске қосу
 # ============================
 iface.launch()