Create app.py

app.py

@@ -0,0 +1,546 @@
+import ast
+import json
+import re
+
+import gradio as gr
+import torch
+from peft import PeftModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+# ---------------------------------------------------------
+# Model configuration
+# ---------------------------------------------------------
+
+BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
+ADAPTER = "mirajbhandari/Entity_Extcation_Quen"
+
+SYSTEM_PROMPT = (
+    "You are an NER model. Extract named entities from the sentence and "
+    'return ONLY a JSON list of objects with keys "text" and "type". '
+    "Allowed types: PERSON, ORGANIZATION, LOCATION, DATE, EVENT, PRODUCT, "
+    "MONEY, TIME, WORK_OF_ART, LANGUAGE, NORP, FAC, GPE."
+)
+
+
+# ---------------------------------------------------------
+# Load tokenizer and model
+# ---------------------------------------------------------
+
+print("Loading tokenizer...")
+
+try:
+    tokenizer = AutoTokenizer.from_pretrained(ADAPTER)
+except Exception:
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+
+if tokenizer.pad_token_id is None:
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+
+
+print("Loading base model...")
+
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL,
+    torch_dtype=dtype,
+    device_map="auto",
+)
+
+
+print("Loading LoRA adapter...")
+
+model = PeftModel.from_pretrained(
+    base_model,
+    ADAPTER,
+)
+
+
+print("Merging LoRA adapter...")
+
+model = model.merge_and_unload()
+model.eval()
+
+print("Model is ready!")
+
+
+# ---------------------------------------------------------
+# Entity colors
+# ---------------------------------------------------------
+
+ENTITY_COLORS = {
+    "PERSON": "#FECACA",
+    "ORGANIZATION": "#BFDBFE",
+    "LOCATION": "#BBF7D0",
+    "GPE": "#A7F3D0",
+    "DATE": "#FDE68A",
+    "TIME": "#FED7AA",
+    "EVENT": "#DDD6FE",
+    "PRODUCT": "#FBCFE8",
+    "MONEY": "#C7D2FE",
+    "WORK_OF_ART": "#E9D5FF",
+    "LANGUAGE": "#BAE6FD",
+    "NORP": "#F5D0FE",
+    "FAC": "#D9F99D",
+}
+
+
+# ---------------------------------------------------------
+# Prompt construction
+# ---------------------------------------------------------
+
+def build_messages(sentence):
+    return [
+        {
+            "role": "system",
+            "content": SYSTEM_PROMPT,
+        },
+        {
+            "role": "user",
+            "content": sentence,
+        },
+    ]
+
+
+# ---------------------------------------------------------
+# JSON parsing
+# ---------------------------------------------------------
+
+def parse_entities(model_output):
+    """
+    Extract and parse the JSON list returned by the model.
+    """
+
+    output = model_output.strip()
+
+    # Remove Markdown code fences if the model adds them.
+    output = re.sub(
+        r"^```(?:json)?\s*",
+        "",
+        output,
+        flags=re.IGNORECASE,
+    )
+
+    output = re.sub(
+        r"\s*```$",
+        "",
+        output,
+    )
+
+    # Locate the JSON list.
+    start = output.find("[")
+    end = output.rfind("]")
+
+    if start == -1 or end == -1 or end < start:
+        raise ValueError(
+            "The model did not return a valid JSON list."
+        )
+
+    json_text = output[start:end + 1]
+
+    try:
+        entities = json.loads(json_text)
+    except json.JSONDecodeError:
+        # Handles occasional Python-style output with single quotes.
+        entities = ast.literal_eval(json_text)
+
+    if not isinstance(entities, list):
+        raise ValueError("NER result must be a list.")
+
+    cleaned_entities = []
+
+    for entity in entities:
+        if not isinstance(entity, dict):
+            continue
+
+        text = str(entity.get("text", "")).strip()
+        entity_type = str(entity.get("type", "")).strip().upper()
+
+        if text and entity_type:
+            cleaned_entities.append(
+                {
+                    "text": text,
+                    "type": entity_type,
+                }
+            )
+
+    return cleaned_entities
+
+
+# ---------------------------------------------------------
+# Find entity positions
+# ---------------------------------------------------------
+
+def find_entity_spans(sentence, entities):
+    """
+    Find the start and end positions of entities in the original sentence.
+
+    This also supports repeated entities.
+    """
+
+    spans = []
+    occupied_positions = []
+
+    for entity in entities:
+        entity_text = entity["text"]
+        entity_type = entity["type"]
+
+        # First try exact matching.
+        matches = list(
+            re.finditer(
+                re.escape(entity_text),
+                sentence,
+            )
+        )
+
+        # If exact matching fails, try case-insensitive matching.
+        if not matches:
+            matches = list(
+                re.finditer(
+                    re.escape(entity_text),
+                    sentence,
+                    flags=re.IGNORECASE,
+                )
+            )
+
+        for match in matches:
+            start = match.start()
+            end = match.end()
+
+            overlaps = any(
+                start < existing_end and end > existing_start
+                for existing_start, existing_end in occupied_positions
+            )
+
+            if overlaps:
+                continue
+
+            spans.append(
+                {
+                    "start": start,
+                    "end": end,
+                    "text": sentence[start:end],
+                    "type": entity_type,
+                }
+            )
+
+            occupied_positions.append((start, end))
+
+            # Use one occurrence for each returned entity.
+            break
+
+    spans.sort(key=lambda item: item["start"])
+
+    return spans
+
+
+# ---------------------------------------------------------
+# Convert spans for Gradio HighlightedText
+# ---------------------------------------------------------
+
+def create_highlighted_output(sentence, spans):
+    """
+    Convert entity spans into the format expected by gr.HighlightedText.
+
+    Example:
+    [
+        ("Barack Obama", "PERSON"),
+        (" visited ", None),
+        ("Paris", "LOCATION")
+    ]
+    """
+
+    if not sentence:
+        return []
+
+    if not spans:
+        return [(sentence, None)]
+
+    highlighted_parts = []
+    current_position = 0
+
+    for span in spans:
+        start = span["start"]
+        end = span["end"]
+
+        if start > current_position:
+            highlighted_parts.append(
+                (sentence[current_position:start], None)
+            )
+
+        highlighted_parts.append(
+            (sentence[start:end], span["type"])
+        )
+
+        current_position = end
+
+    if current_position < len(sentence):
+        highlighted_parts.append(
+            (sentence[current_position:], None)
+        )
+
+    return highlighted_parts
+
+
+# ---------------------------------------------------------
+# Model inference
+# ---------------------------------------------------------
+
+@torch.inference_mode()
+def extract_entities(sentence):
+    sentence = sentence.strip()
+
+    if not sentence:
+        raise gr.Error("Please enter a sentence.")
+
+    messages = build_messages(sentence)
+
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+    )
+
+    model_inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+    )
+
+    model_device = next(model.parameters()).device
+
+    model_inputs = {
+        key: value.to(model_device)
+        for key, value in model_inputs.items()
+    }
+
+    generated_ids = model.generate(
+        **model_inputs,
+        max_new_tokens=256,
+        do_sample=False,
+        repetition_penalty=1.05,
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
+    )
+
+    # Remove the original prompt tokens.
+    generated_tokens = generated_ids[
+        :,
+        model_inputs["input_ids"].shape[1]:
+    ]
+
+    model_output = tokenizer.batch_decode(
+        generated_tokens,
+        skip_special_tokens=True,
+    )[0].strip()
+
+    try:
+        entities = parse_entities(model_output)
+    except Exception as error:
+        return (
+            [(sentence, None)],
+            [],
+            {
+                "error": str(error),
+                "raw_model_output": model_output,
+            },
+        )
+
+    spans = find_entity_spans(sentence, entities)
+    highlighted_output = create_highlighted_output(sentence, spans)
+
+    entity_table = [
+        [
+            span["text"],
+            span["type"],
+            span["start"],
+            span["end"],
+        ]
+        for span in spans
+    ]
+
+    json_output = {
+        "sentence": sentence,
+        "entities": [
+            {
+                "text": span["text"],
+                "type": span["type"],
+                "start": span["start"],
+                "end": span["end"],
+            }
+            for span in spans
+        ],
+    }
+
+    return highlighted_output, entity_table, json_output
+
+
+# ---------------------------------------------------------
+# Clear interface
+# ---------------------------------------------------------
+
+def clear_outputs():
+    return "", [], [], None
+
+
+# ---------------------------------------------------------
+# Gradio user interface
+# ---------------------------------------------------------
+
+CUSTOM_CSS = """
+.gradio-container {
+    max-width: 1100px !important;
+    margin: auto !important;
+}
+
+#main-title {
+    text-align: center;
+    margin-bottom: 4px;
+}
+
+#subtitle {
+    text-align: center;
+    color: #64748b;
+    margin-bottom: 24px;
+}
+
+#input-card,
+#result-card {
+    border-radius: 14px;
+}
+"""
+
+
+with gr.Blocks(
+    title="Named Entity Recognition",
+    css=CUSTOM_CSS,
+    theme=gr.themes.Soft(),
+) as demo:
+
+    gr.Markdown(
+        """
+        # Named Entity Recognition
+        """,
+        elem_id="main-title",
+    )
+
+    gr.Markdown(
+        """
+        Enter a sentence to identify and highlight named entities.
+        """,
+        elem_id="subtitle",
+    )
+
+    with gr.Group(elem_id="input-card"):
+        sentence_input = gr.Textbox(
+            label="Original sentence",
+            placeholder=(
+                "Example: Sundar Pichai visited Google headquarters "
+                "in California on July 15, 2026."
+            ),
+            lines=4,
+        )
+
+        with gr.Row():
+            extract_button = gr.Button(
+                "Extract Entities",
+                variant="primary",
+            )
+
+            clear_button = gr.Button(
+                "Clear",
+                variant="secondary",
+            )
+
+    with gr.Group(elem_id="result-card"):
+        highlighted_output = gr.HighlightedText(
+            label="Highlighted sentence",
+            color_map=ENTITY_COLORS,
+            show_legend=True,
+            show_inline_category=True,
+            combine_adjacent=True,
+        )
+
+        entity_table = gr.Dataframe(
+            headers=[
+                "Entity",
+                "Entity Type",
+                "Start Position",
+                "End Position",
+            ],
+            datatype=[
+                "str",
+                "str",
+                "number",
+                "number",
+            ],
+            label="Detected entities",
+            interactive=False,
+            wrap=True,
+        )
+
+        with gr.Accordion(
+            "JSON output",
+            open=False,
+        ):
+            json_output = gr.JSON(
+                label="Structured NER result"
+            )
+
+    gr.Examples(
+        examples=[
+            [
+                "Sundar Pichai is the CEO of Google and lives in California."
+            ],
+            [
+                "Apple launched the iPhone in September 2025."
+            ],
+            [
+                "Barack Obama visited Paris on January 10, 2024."
+            ],
+            [
+                "Microsoft invested 10 billion dollars in OpenAI."
+            ],
+            [
+                "The FIFA World Cup was held in Qatar in 2022."
+            ],
+        ],
+        inputs=sentence_input,
+    )
+
+    extract_button.click(
+        fn=extract_entities,
+        inputs=sentence_input,
+        outputs=[
+            highlighted_output,
+            entity_table,
+            json_output,
+        ],
+        api_name="extract_entities",
+    )
+
+    sentence_input.submit(
+        fn=extract_entities,
+        inputs=sentence_input,
+        outputs=[
+            highlighted_output,
+            entity_table,
+            json_output,
+        ],
+    )
+
+    clear_button.click(
+        fn=clear_outputs,
+        inputs=[],
+        outputs=[
+            sentence_input,
+            highlighted_output,
+            entity_table,
+            json_output,
+        ],
+        queue=False,
+    )
+
+
+if __name__ == "__main__":
+    demo.queue().launch()