feat: update app

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from huggingface_hub import HfApi
+from huggingface_hub import HfApi, create_repo
 import os
 import re
 import json
@@ -8,6 +8,11 @@ import random
 from typing import List, Dict, Union, Tuple
 from gliner import GLiNER
 from datasets import load_dataset
+from dotenv import load_dotenv
+
+# Load environment variables from .env
+load_dotenv()
+HF_TOKEN = os.getenv("HUGGINGFACE_ACCESS_TOKEN")
 
 # Available models for annotation
 AVAILABLE_MODELS = [
@@ -138,8 +143,58 @@ def extract_tokens_and_labels(data: List[Dict[str, Union[str, None]]]) -> Dict[s
 # Global variables for dataset viewer
 dynamic_dataset = None
 
+def load_dataset():
+    global dynamic_dataset
+    try:
+        with open("data/annotated_data.json", 'rt') as dataset:
+            ANNOTATED_DATA = json.load(dataset)
+        dynamic_dataset = DynamicDataset(ANNOTATED_DATA)
+        max_value = len(dynamic_dataset.data) - 1 if dynamic_dataset.data else 0
+        return prepare_for_highlight(dynamic_dataset.load_current_example()), gr.update(value=0, maximum=max_value)
+    except Exception as e:
+        return [("Error loading dataset: " + str(e), None)], gr.update(value=0, maximum=1)
+
+def example_by_id(id):
+    global dynamic_dataset
+    if dynamic_dataset is None:
+        return [("Please load a dataset first", None)], gr.update(value=0, maximum=1)
+    try:
+        id = int(id)  # Ensure id is an integer
+        dynamic_dataset.example_by_id(id)
+        current = dynamic_dataset.current
+        max_value = len(dynamic_dataset.data) - 1
+        return prepare_for_highlight(dynamic_dataset.load_current_example()), gr.update(value=current, maximum=max_value)
+    except Exception as e:
+        return [("Error navigating to example: " + str(e), None)], gr.update(value=0, maximum=1)
+
+def next_example():
+    global dynamic_dataset
+    if dynamic_dataset is None:
+        return [("Please load a dataset first", None)], gr.update(value=0, maximum=1)
+    try:
+        dynamic_dataset.next_example()
+        current = dynamic_dataset.current
+        max_value = len(dynamic_dataset.data) - 1
+        return prepare_for_highlight(dynamic_dataset.load_current_example()), gr.update(value=current, maximum=max_value)
+    except Exception as e:
+        return [("Error navigating to next example: " + str(e), None)], gr.update(value=0, maximum=1)
+
+def previous_example():
+    global dynamic_dataset
+    if dynamic_dataset is None:
+        return [("Please load a dataset first", None)], gr.update(value=0, maximum=1)
+    try:
+        dynamic_dataset.previous_example()
+        current = dynamic_dataset.current
+        max_value = len(dynamic_dataset.data) - 1
+        return prepare_for_highlight(dynamic_dataset.load_current_example()), gr.update(value=current, maximum=max_value)
+    except Exception as e:
+        return [("Error navigating to previous example: " + str(e), None)], gr.update(value=0, maximum=1)
+
 def update_example(data):
     global dynamic_dataset
+    if dynamic_dataset is None:
+        return [("Please load a dataset first", None)]
     tokens, ner = extract_tokens_and_labels(data)
     dynamic_dataset.data[dynamic_dataset.current]["tokenized_text"] = tokens
     dynamic_dataset.data[dynamic_dataset.current]["ner"] = ner
@@ -147,36 +202,19 @@ def update_example(data):
 
 def validate_example():
     global dynamic_dataset
+    if dynamic_dataset is None:
+        return [("Please load a dataset first", None)]
     dynamic_dataset.data[dynamic_dataset.current]["validated"] = True
     return [("The example was validated!", None)]
 
-def next_example():
-    global dynamic_dataset
-    dynamic_dataset.next_example()
-    return prepare_for_highlight(dynamic_dataset.load_current_example()), dynamic_dataset.current
-
-def previous_example():
-    global dynamic_dataset
-    dynamic_dataset.previous_example()
-    return prepare_for_highlight(dynamic_dataset.load_current_example()), dynamic_dataset.current
-
 def save_dataset(inp):
     global dynamic_dataset
+    if dynamic_dataset is None:
+        return [("Please load a dataset first", None)]
     with open("data/annotated_data.json", "wt") as file:
         json.dump(dynamic_dataset.data, file)
     return [("The validated dataset was saved as data/annotated_data.json", None)]
 
-def load_dataset():
-    global dynamic_dataset
-    try:
-        with open("data/annotated_data.json", 'rt') as dataset:
-            ANNOTATED_DATA = json.load(dataset)
-        dynamic_dataset = DynamicDataset(ANNOTATED_DATA)
-        max_value = len(dynamic_dataset.data) - 1 if dynamic_dataset.data else 0
-        return prepare_for_highlight(dynamic_dataset.load_current_example()), 0, max_value
-    except Exception as e:
-        return [("Error loading dataset: " + str(e), None)], 0, 0
-
 # Original annotation functions
 def transform_data(data):
     tokens = tokenize_text(data['text'])
@@ -209,8 +247,9 @@ def merge_entities(entities):
     merged.append(current)
     return merged
 
-def annotate_text(model, text, labels: List[str], threshold: float, nested_ner: bool) -> Dict:
+def annotate_text(model: GLiNER, text, labels: List[str], threshold: float, nested_ner: bool) -> Dict:
     labels = [label.strip() for label in labels]
+    entities = model.predict_entities(text, labels, flat_ner=not nested_ner, threshold=threshold)
     r = {
         "text": text,
         "entities": [
@@ -221,14 +260,36 @@ def annotate_text(model, text, labels: List[str], threshold: float, nested_ner:
                 "end": entity["end"],
                 "score": 0,
             }
-            for entity in model.predict_entities(
-                text, labels, flat_ner=not nested_ner, threshold=threshold
-            )
+            for entity in entities
         ],
     }
     r["entities"] = merge_entities(r["entities"])
     return transform_data(r)
 
+def batch_annotate_text(model: GLiNER, texts: List[str], labels: List[str], threshold: float, nested_ner: bool) -> List[Dict]:
+    """Annotate multiple texts in batch"""
+    labels = [label.strip() for label in labels]
+    batch_entities = model.batch_predict_entities(texts, labels, flat_ner=not nested_ner, threshold=threshold)
+    
+    results = []
+    for text, entities in zip(texts, batch_entities):
+        r = {
+            "text": text,
+            "entities": [
+                {
+                    "entity": entity["label"],
+                    "word": entity["text"],
+                    "start": entity["start"],
+                    "end": entity["end"],
+                    "score": 0,
+                }
+                for entity in entities
+            ],
+        }
+        r["entities"] = merge_entities(r["entities"])
+        results.append(transform_data(r))
+    return results
+
 class AutoAnnotator:
     def __init__(
         self, model: str = "knowledgator/gliner-multitask-large-v0.5",
@@ -248,20 +309,29 @@ class AutoAnnotator:
             ) -> List[Dict]:
         self.stat["total"] = len(data)
         self.stat["current"] = -1  # Reset current progress
-        for text in data:
-            self.stat["current"] += 1
+        
+        # Process texts in batches
+        batch_size = 32  # Adjust based on your GPU memory
+        processed_data = []
+        
+        for i in range(0, len(data), batch_size):
+            batch_texts = data[i:i + batch_size]
             if isinstance(prompt, list):
                 prompt_text = random.choice(prompt)
             else:
                 prompt_text = prompt
-            text = f"{prompt_text}\n{text}" if prompt_text else text
-
-            annotation = annotate_text(self.model, text, labels, threshold, nested_ner)
-
-            if not annotation["ner"]:  # If no entities identified
-                annotation = {"tokenized_text": tokenize_text(text), "ner": [], "validated": False}
-
-            self.annotated_data.append(annotation)
+            
+            # Add prompt to each text in batch
+            batch_texts = [f"{prompt_text}\n{text}" if prompt_text else text for text in batch_texts]
+            
+            # Process batch
+            batch_results = batch_annotate_text(self.model, batch_texts, labels, threshold, nested_ner)
+            processed_data.extend(batch_results)
+            
+            # Update progress
+            self.stat["current"] = min(i + batch_size, len(data))
+        
+        self.annotated_data = processed_data
         return self.annotated_data
 
 # Global variables
@@ -281,31 +351,61 @@ def process_uploaded_file(file_obj):
     except Exception as e:
         return f"Error reading file: {str(e)}"
 
-def annotate(model, labels, threshold, prompt):
+def is_valid_repo_name(repo_name):
+    # Hugging Face repo names must not contain slashes or spaces
+    return bool(re.match(r'^[A-Za-z0-9_.-]+$', repo_name))
+
+def create_hf_repo(repo_name: str, repo_type: str = "dataset", private: bool = False):
+    """Create a new repository on Hugging Face Hub"""
+    if not is_valid_repo_name(repo_name):
+        raise Exception("Invalid repo name: must not contain slashes, spaces, or special characters except '-', '_', '.'")
+    try:
+        api = HfApi(token=HF_TOKEN)
+        user = api.whoami()['name']
+        repo_id = f"{user}/{repo_name}"
+        create_repo(
+            repo_id=repo_id,
+            repo_type=repo_type,
+            private=private,
+            exist_ok=True,
+            token=HF_TOKEN
+        )
+        return repo_id
+    except Exception as e:
+        raise Exception(f"Error creating repository: {str(e)}")
+
+def annotate(model, labels, threshold, prompt, save_to_hub, repo_name, repo_type, is_private):
     global annotator
     try:
         if not sentences:
             return "Please upload a file with text first!"
-            
+        if save_to_hub and not is_valid_repo_name(repo_name):
+            return "Error: Invalid repo name. Only use letters, numbers, '-', '_', or '.' (no slashes or spaces)."
         labels = [label.strip() for label in labels.split(",")]
         annotator = AutoAnnotator(model)
         annotated_data = annotator.auto_annotate(sentences, labels, prompt, threshold)
-        
-        # Save annotated data
+        # Save annotated data locally
         os.makedirs("data", exist_ok=True)
-        with open("data/annotated_data.json", "wt") as file:
+        local_path = "data/annotated_data.json"
+        with open(local_path, "wt") as file:
             json.dump(annotated_data, file, ensure_ascii=False)
-            
-        # Upload to Hugging Face Hub
-        api = HfApi()
-        api.upload_file(
-            path_or_fileobj="data/annotated_data.json",
-            path_in_repo="annotated_data.json",
-            repo_id="YOUR_USERNAME/YOUR_REPO_NAME",  # Replace with your repo
-            repo_type="dataset"
-        )
-        
-        return "Successfully annotated and saved to Hugging Face Hub!"
+        status_messages = [f"Successfully annotated and saved locally to {local_path}"]
+        # Upload to Hugging Face Hub if requested
+        if save_to_hub:
+            try:
+                repo_id = create_hf_repo(repo_name, repo_type, is_private)
+                api = HfApi(token=HF_TOKEN)
+                api.upload_file(
+                    path_or_fileobj=local_path,
+                    path_in_repo="annotated_data.json",
+                    repo_id=repo_id,
+                    repo_type=repo_type,
+                    token=HF_TOKEN
+                )
+                status_messages.append(f"Successfully uploaded to Hugging Face Hub repository: {repo_id}")
+            except Exception as e:
+                status_messages.append(f"Error with Hugging Face Hub: {str(e)}")
+        return "\n".join(status_messages)
     except Exception as e:
         return f"Error during annotation: {str(e)}"
 
@@ -504,6 +604,76 @@ def process_local_file(file_obj, file_format):
     except Exception as e:
         return f"Error processing file: {str(e)}"
 
+# Add a function to download the annotated data
+
+def download_annotated_data():
+    file_path = "data/annotated_data.json"
+    if os.path.exists(file_path):
+        return file_path
+    else:
+        return None
+
+def download_to_folder():
+    """Download annotated data to a local folder"""
+    try:
+        source_path = "data/annotated_data.json"
+        if not os.path.exists(source_path):
+            return "No annotated data found!"
+        
+        # Create downloads directory if it doesn't exist
+        download_dir = os.path.expanduser("~/Downloads")
+        os.makedirs(download_dir, exist_ok=True)
+        
+        # Copy file to downloads folder
+        import shutil
+        dest_path = os.path.join(download_dir, "annotated_data.json")
+        shutil.copy2(source_path, dest_path)
+        return f"Successfully downloaded to {dest_path}"
+    except Exception as e:
+        return f"Error downloading file: {str(e)}"
+
+def update_hf_dataset(repo_name: str, repo_type: str = "dataset", is_private: bool = False):
+    """Update or create a Hugging Face dataset with the current annotated data"""
+    try:
+        if not dynamic_dataset or not dynamic_dataset.data:
+            return "No data to upload! Please load or annotate data first."
+        
+        # Save current data to local file
+        os.makedirs("data", exist_ok=True)
+        local_path = "data/annotated_data.json"
+        with open(local_path, "wt") as file:
+            json.dump(dynamic_dataset.data, file, ensure_ascii=False)
+        
+        # Create or update repository
+        try:
+            repo_id = create_hf_repo(repo_name, repo_type, is_private)
+            api = HfApi(token=HF_TOKEN)
+            api.upload_file(
+                path_or_fileobj=local_path,
+                path_in_repo="annotated_data.json",
+                repo_id=repo_id,
+                repo_type=repo_type,
+                token=HF_TOKEN
+            )
+            return f"Successfully uploaded to Hugging Face Hub repository: {repo_id}"
+        except Exception as e:
+            if "already exists" in str(e):
+                # If repo exists, just update the file
+                user = api.whoami()['name']
+                repo_id = f"{user}/{repo_name}"
+                api.upload_file(
+                    path_or_fileobj=local_path,
+                    path_in_repo="annotated_data.json",
+                    repo_id=repo_id,
+                    repo_type=repo_type,
+                    token=HF_TOKEN
+                )
+                return f"Successfully updated existing repository: {repo_id}"
+            else:
+                raise e
+    except Exception as e:
+        return f"Error updating Hugging Face dataset: {str(e)}"
+
 # Create the main interface with tabs
 with gr.Blocks() as demo:
     gr.Markdown("# NER Annotation Tool")
@@ -539,14 +709,77 @@ with gr.Blocks() as demo:
                         placeholder="Enter your annotation prompt (optional)",
                         scale=2
                     )
+                    
+                    with gr.Group():
+                        gr.Markdown("### Save Options")
+                        save_to_hub = gr.Checkbox(
+                            label="Save to Hugging Face Hub",
+                            value=False
+                        )
+                        
+                        with gr.Group(visible=False) as hub_settings:
+                            gr.Markdown("#### Hugging Face Hub Settings")
+                            repo_name = gr.Textbox(
+                                label="Repository Name",
+                                placeholder="Enter repository name (e.g., my-ner-dataset)",
+                                scale=2
+                            )
+                            repo_type = gr.Dropdown(
+                                choices=["dataset", "model", "space"],
+                                value="dataset",
+                                label="Repository Type"
+                            )
+                            is_private = gr.Checkbox(
+                                label="Private Repository",
+                                value=False
+                            )
+                    
                     annotate_btn = gr.Button("Annotate Data")
                     output_info = gr.Textbox(label="Processing Status")
                     
+                    # Add download buttons for annotated data
+                    with gr.Row():
+                        download_btn_annot = gr.Button("Download Annotated Data", visible=False)
+                    download_file_annot = gr.File(label="Download", interactive=False, visible=False)
+                    download_status = gr.Textbox(label="Download Status", visible=False)
+                    
+                    def toggle_hub_settings(save_to_hub):
+                        return {
+                            hub_settings: gr.update(visible=save_to_hub)
+                        }
+                    
+                    save_to_hub.change(
+                        fn=toggle_hub_settings,
+                        inputs=[save_to_hub],
+                        outputs=[hub_settings]
+                    )
+                    
+                    def show_download_buttons(status):
+                        # Show download buttons only if annotation was successful
+                        if status and status.startswith("Successfully annotated and saved locally"):
+                            return gr.update(visible=True), gr.update(visible=True)
+                        return gr.update(visible=False), gr.update(visible=False)
+                    
                     annotate_btn.click(
                         fn=annotate,
-                        inputs=[model, labels, threshold, prompt],
+                        inputs=[
+                            model, labels, threshold, prompt,
+                            save_to_hub, repo_name, repo_type, is_private
+                        ],
                         outputs=[output_info]
                     )
+                    output_info.change(
+                        fn=show_download_buttons,
+                        inputs=[output_info],
+                        outputs=[download_btn_annot, download_status]
+                    )
+                    def handle_download_annot():
+                        file_path = download_annotated_data()
+                        if file_path:
+                            return gr.update(value=file_path, visible=True)
+                        else:
+                            return gr.update(visible=False)
+                    download_btn_annot.click(fn=handle_download_annot, inputs=None, outputs=[download_file_annot])
         
         with gr.TabItem("Dataset Viewer"):
             with gr.Row():
@@ -576,7 +809,14 @@ with gr.Blocks() as demo:
                         visible=False
                     )
                     
-                    bar = gr.Slider(minimum=0, maximum=1, step=1, label="Progress", interactive=False)
+                    bar = gr.Slider(
+                        minimum=0,
+                        maximum=1,
+                        step=1,
+                        label="Progress",
+                        interactive=True,
+                        info="Use slider to navigate through examples"
+                    )
                     
                     with gr.Row():
                         previous_btn = gr.Button("Previous example")
@@ -586,6 +826,26 @@ with gr.Blocks() as demo:
                     validate_btn = gr.Button("Validate")
                     save_btn = gr.Button("Save validated dataset")
                     
+                    # Add Hugging Face upload section
+                    with gr.Group():
+                        gr.Markdown("### Upload to Hugging Face")
+                        hf_repo_name = gr.Textbox(
+                            label="Repository Name",
+                            placeholder="Enter repository name (e.g., my-ner-dataset)",
+                            scale=2
+                        )
+                        hf_repo_type = gr.Dropdown(
+                            choices=["dataset", "model", "space"],
+                            value="dataset",
+                            label="Repository Type"
+                        )
+                        hf_is_private = gr.Checkbox(
+                            label="Private Repository",
+                            value=False
+                        )
+                        upload_to_hf_btn = gr.Button("Upload to Hugging Face")
+                        hf_upload_status = gr.Textbox(label="Upload Status")
+                    
                     inp_box = gr.HighlightedText(value=None, interactive=True)
                     
                     def toggle_local_inputs():
@@ -647,5 +907,18 @@ with gr.Blocks() as demo:
                     validate_btn.click(fn=validate_example, inputs=None, outputs=inp_box)
                     next_btn.click(fn=next_example, inputs=None, outputs=[inp_box, bar])
                     previous_btn.click(fn=previous_example, inputs=None, outputs=[inp_box, bar])
+                    bar.change(
+                        fn=example_by_id,
+                        inputs=[bar],
+                        outputs=[inp_box, bar],
+                        api_name="example_by_id"
+                    )
+                    
+                    # Add Hugging Face upload functionality
+                    upload_to_hf_btn.click(
+                        fn=update_hf_dataset,
+                        inputs=[hf_repo_name, hf_repo_type, hf_is_private],
+                        outputs=[hf_upload_status]
+                    )
 
 demo.launch()

pyproject.toml

@@ -9,4 +9,5 @@ dependencies = [
     "gliner>=0.2.20",
     "gradio>=5.31.0",
     "huggingface-hub>=0.32.1",
+    "python-dotenv>=1.1.0",
 ]

requirements.txt

@@ -1,4 +1,5 @@
 gradio==5.31.0
 datasets>=3.6.0
 gliner>=0.2.20
-huggingface-hub>=0.32.1
+huggingface-hub>=0.32.1
+python-dotenv>=1.1.0

uv.lock

@@ -902,6 +902,7 @@ dependencies = [
     { name = "gliner" },
     { name = "gradio" },
     { name = "huggingface-hub" },
+    { name = "python-dotenv" },
 ]
 
 [package.metadata]
@@ -910,6 +911,7 @@ requires-dist = [
     { name = "gliner", specifier = ">=0.2.20" },
     { name = "gradio", specifier = ">=5.31.0" },
     { name = "huggingface-hub", specifier = ">=0.32.1" },
+    { name = "python-dotenv", specifier = ">=1.1.0" },
 ]
 
 [[package]]
@@ -1646,6 +1648,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" },
 ]
 
+[[package]]
+name = "python-dotenv"
+version = "1.1.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/88/2c/7bb1416c5620485aa793f2de31d3df393d3686aa8a8506d11e10e13c5baf/python_dotenv-1.1.0.tar.gz", hash = "sha256:41f90bc6f5f177fb41f53e87666db362025010eb28f60a01c9143bfa33a2b2d5", size = 39920, upload-time = "2025-03-25T10:14:56.835Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/1e/18/98a99ad95133c6a6e2005fe89faedf294a748bd5dc803008059409ac9b1e/python_dotenv-1.1.0-py3-none-any.whl", hash = "sha256:d7c01d9e2293916c18baf562d95698754b0dbbb5e74d457c45d4f6561fb9d55d", size = 20256, upload-time = "2025-03-25T10:14:55.034Z" },
+]
+
 [[package]]
 name = "python-multipart"
 version = "0.0.20"