Update app.py

app.py

@@ -1,50 +1,74 @@
+import spaces
+
 import gradio as gr
-from transformers import pipeline
+import torch
+from transformers import AutoTokenizer, AutoModelForTokenClassification
 
 MODEL_NAME = "122Uswa/code-switching-codesaviours-si26-uswa"
 
-classifier = pipeline(
-    "token-classification",
-    model=MODEL_NAME,
-    aggregation_strategy="simple"
-)
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForTokenClassification.from_pretrained(MODEL_NAME)
+
+model.eval()
 
+
+@spaces.GPU(duration=60)
 def predict_language(text):
     if not text.strip():
         return "Please enter a sentence."
 
     words = text.split()
-    results = classifier(text)
 
-    output = []
+    # Tokenize words separately so word_ids() can map predictions
+    inputs = tokenizer(
+        words,
+        is_split_into_words=True,
+        return_tensors="pt",
+        truncation=True
+    )
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    predictions = torch.argmax(outputs.logits, dim=2)[0].tolist()
+
+    word_ids = tokenizer(
+        words,
+        is_split_into_words=True,
+        truncation=True
+    ).word_ids()
+
+    results = []
+    previous_word = None
+
+    for prediction, word_id in zip(predictions, word_ids):
+
+        if word_id is None or word_id == previous_word:
+            continue
 
-    for word in words:
-        label = "UNKNOWN"
+        word = words[word_id]
+        label = model.config.id2label[prediction]
 
-        # Find prediction overlapping with this word
-        for result in results:
-            predicted_word = result["word"].strip()
+        results.append(f"{word} → {label}")
 
-            if predicted_word.lower() == word.lower():
-                label = result["entity_group"]
-                break
+        previous_word = word_id
 
-        output.append(f"{word} → {label}")
+    return "\n".join(results)
 
-    return "\n".join(output)
 
 demo = gr.Interface(
     fn=predict_language,
     inputs=gr.Textbox(
         label="Roman Urdu / English Sentence",
-        placeholder="Yaar mujhe time chahiye for this assignment",
+        placeholder="Yaar mujhe thora time chahiye for this assignment",
         lines=3
     ),
     outputs=gr.Textbox(
         label="Language Identification"
     ),
     title="Roman Urdu–English Language Identifier",
-    description="Identifies each word as Urdu (URD), English (ENG), or Mixed (MIX)."
+    description="Identifies each word as URD, ENG, or MIX."
 )
 
 demo.launch()