Updated Gradio App

.github/workflows/ci-cd.yml

@@ -85,7 +85,7 @@ jobs:
           git merge --no-ff origin/develop
           git push origin main
 
-  sync-develop:
+  finalize-deployment:
     needs: merge-to-main
     runs-on: ubuntu-latest
     steps:

app.py

@@ -1,12 +1,10 @@
 import gradio as gr
 from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification
 
-# Load fine-tuned model
-model_path = "./ner_model"
+model_path = "./models/ner_model"
 tokenizer = AutoTokenizer.from_pretrained(model_path)
 model = AutoModelForTokenClassification.from_pretrained(model_path)
 
-# Create NER pipeline
 ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
 
 def ner_prediction(text):

train.py

@@ -1,3 +1,4 @@
+import os
 import torch
 from transformers import AutoTokenizer, AutoModelForTokenClassification, TrainingArguments, Trainer
 from datasets import load_dataset, load_metric
@@ -21,7 +22,7 @@ model = AutoModelForTokenClassification.from_pretrained(model_checkpoint, num_la
 
 # Training arguments
 training_args = TrainingArguments(
-    output_dir="./ner_model",
+    output_dir="./models/ner_model",
     evaluation_strategy="epoch",
     save_strategy="epoch",
     learning_rate=2e-5,
@@ -51,6 +52,10 @@ trainer = Trainer(
 # Train model
 trainer.train()
 
+# Ensure directory exists before saving
+output_dir = "./models/ner_model"
+os.makedirs(output_dir, exist_ok=True)
+
 # Save model
-trainer.save_model("./ner_model")
-tokenizer.save_pretrained("./ner_model")
+trainer.save_model(output_dir)
+tokenizer.save_pretrained(output_dir)