Update app.py

app.py

@@ -4,86 +4,139 @@ import json
 import uuid
 import torch
 import datetime
-import pandas as pd
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-from huggingface_hub import HfApi, create_repo, upload_file
-from datasets import Dataset
+import torch.nn as nn
+from transformers import AutoTokenizer, AutoModel, AutoConfig
+from huggingface_hub import HfApi, create_repo, hf_hub_download
+from torchcrf import CRF
 
-# Configuration
-MODEL_NAME = "distilbert-base-uncased-finetuned-sst-2-english"
+# Constants
 HF_DATASET_REPO = "M2ai/mgtd-logs"
 HF_TOKEN = os.getenv("Mgtd")
 DATASET_CREATED = False
 
-# Load model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
-
-# Make directories
-os.makedirs("logs", exist_ok=True)
-
+# Model identifiers
+code = "ENG"
+pntr = 2
+model_name_or_path = "microsoft/mdeberta-v3-base"
+hf_token = os.environ.get("HF_WRITE")  # Set this before running
+
+# Download model checkpoint
+file_path = hf_hub_download(
+    repo_id="1024m/MGTD-Long-New",
+    filename=f"{code}/mdeberta-epoch-{pntr}.pt",
+    token=hf_token,
+    local_dir="./checkpoints"
+)
+
+# Define CRF model
+class AutoModelCRF(nn.Module):
+    def __init__(self, model_name_or_path, dropout=0.075):
+        super().__init__()
+        self.config = AutoConfig.from_pretrained(model_name_or_path)
+        self.num_labels = 2
+        self.encoder = AutoModel.from_pretrained(model_name_or_path, trust_remote_code=True, config=self.config)
+        self.dropout = nn.Dropout(dropout)
+        self.linear = nn.Linear(self.config.hidden_size, self.num_labels)
+        self.crf = CRF(self.num_labels, batch_first=True)
+
+    def forward(self, input_ids, attention_mask):
+        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
+        seq_output = self.dropout(outputs[0])
+        emissions = self.linear(seq_output)
+        tags = self.crf.decode(emissions, attention_mask.byte())
+        return tags, emissions
+
+# Load model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+model = AutoModelCRF(model_name_or_path)
+checkpoint = torch.load(file_path, map_location="cpu")
+model.load_state_dict(checkpoint.get("model_state_dict", checkpoint), strict=False)
+model = model.to(device)
+model.eval()
+
+# Inference function
+def get_word_classifications(text):
+    text = " ".join(text.split(" ")[:2048])
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    inputs = {k: v.to(device) for k, v in inputs.items()}
+    tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
+    with torch.no_grad():
+        tags, _ = model(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"])
+    word_tags = []
+    current_word = ""
+    current_tag = None
+    for token, tag in zip(tokens, tags[0]):
+        if token in ["<s>", "</s>"]:
+            continue
+        if token.startswith("▁"):
+            if current_word:
+                word_tags.append(str(current_tag))
+            current_word = token[1:] if token != "▁" else ""
+            current_tag = tag
+        else:
+            current_word += token
+    if current_word:
+        word_tags.append(str(current_tag))
+    return word_tags
+
+# HF logging setup
 def setup_hf_dataset():
     global DATASET_CREATED
     if not DATASET_CREATED and HF_TOKEN:
         try:
-            api = HfApi()
             create_repo(HF_DATASET_REPO, repo_type="dataset", token=HF_TOKEN, exist_ok=True)
             DATASET_CREATED = True
-            print(f"Dataset {HF_DATASET_REPO} is ready")
+            print(f"Dataset {HF_DATASET_REPO} is ready.")
         except Exception as e:
             print(f"Error setting up dataset: {e}")
-    elif not HF_TOKEN:
-        print("Warning: HF_TOKEN not set. Logs will be saved locally only.")
 
+# Main inference + logging function
 def infer_and_log(text_input):
-    inputs = tokenizer(text_input, return_tensors="pt", truncation=True)
-    with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits.tolist()
-        predicted = torch.argmax(outputs.logits, dim=-1).item()
-        label = model.config.id2label[predicted]
-
+    word_tags = get_word_classifications(text_input)
     timestamp = datetime.datetime.now().isoformat()
     submission_id = str(uuid.uuid4())
+
     log_data = {
         "id": submission_id,
         "timestamp": timestamp,
         "input": text_input,
-        "logits": logits
+        "output_tags": word_tags
     }
 
+    os.makedirs("logs", exist_ok=True)
     log_file = f"logs/{timestamp.replace(':', '_')}.json"
     with open(log_file, "w") as f:
         json.dump(log_data, f, indent=2)
 
     if HF_TOKEN and DATASET_CREATED:
         try:
-            api = HfApi()
-            api.upload_file(
+            HfApi().upload_file(
                 path_or_fileobj=log_file,
                 path_in_repo=f"logs/{os.path.basename(log_file)}",
                 repo_id=HF_DATASET_REPO,
                 repo_type="dataset",
                 token=HF_TOKEN
             )
-            print(f"Uploaded log {submission_id} to {HF_DATASET_REPO}")
+            print(f"Uploaded log {submission_id}")
         except Exception as e:
-            print(f"Error uploading to HF dataset: {e}")
+            print(f"Error uploading log: {e}")
 
-    return label
+    return " ".join(word_tags)
 
 def clear_fields():
     return "", ""
 
-# Setup the dataset on startup
+# Prepare dataset once
 setup_hf_dataset()
 
+# Gradio UI
 with gr.Blocks() as app:
-    gr.Markdown("## AI Text Detector")
+    gr.Markdown("## MDeBERTa+CRF Word Tagger")
 
     with gr.Row():
-        input_box = gr.Textbox(label="Input Text", lines=10, interactive=True)
-        output_box = gr.Textbox(label="Output", lines=2, interactive=False)
+        input_box = gr.Textbox(label="Input Text", lines=10)
+        output_box = gr.Textbox(label="Output Tags", lines=2, interactive=False)
 
     with gr.Row():
         submit_btn = gr.Button("Submit")