Upload handler.py with huggingface_hub

handler.py

@@ -5,7 +5,7 @@ import torch
 import torch.nn as nn
 import re
 
-# 1. Define Architecture Patch
+# 1. Architecture Patch (RMSNorm)
 class RMSNorm(nn.Module):
     def __init__(self, dim: int, eps: float = 1e-6):
         super().__init__()
@@ -25,36 +25,45 @@ def replace_layernorm_with_rmsnorm(module: nn.Module):
         else:
             replace_layernorm_with_rmsnorm(child)
 
-# 2. Define The Glue Logic (UPDATED)
+# 2. Glue Logic (Prefixes + Punctuation)
 def fix_arabic_output(text):
     if not text: return text
-    # Glue Prefixes (Next word)
+    
+    # A. Glue Prefixes (Al, Lil, Wa-Al, Bi-Al)
     prefix_pattern = r'(^|\s)(ال|لل|وال|بال)\s+(?=\S)'
     text = re.sub(prefix_pattern, r'\1\2', text)
-    text = re.sub(prefix_pattern, r'\1\2', text)
-    # Glue Punctuation (Previous word)
-    punctuation_pattern = r'\s+([،؟!.,])'
-    text = re.sub(punctuation_pattern, r'\1', text)
+    text = re.sub(prefix_pattern, r'\1\2', text) # Twice for safety
+    
+    # B. Glue Punctuation (Remove space before punctuation)
+    punctuation_marks = r'[،؟!.,]'
+    text = re.sub(r'\s+(' + punctuation_marks + ')', r'\1', text)
+    
     return text.strip()
 
 class EndpointHandler:
     def __init__(self, path=""):
+        # Load Config & Skeleton
         config = AutoConfig.from_pretrained(path)
         self.model = AutoModelForSeq2SeqLM.from_config(config)
         replace_layernorm_with_rmsnorm(self.model)
         
+        # Load Weights Safely
         try:
+            # Try standard load first
+            self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
+            replace_layernorm_with_rmsnorm(self.model)
+        except:
+             # Fallback: Load state dict manually with strict=False
              from safetensors.torch import load_file
              import os
              w_path = os.path.join(path, "model.safetensors")
              if os.path.exists(w_path):
-                 self.model.load_state_dict(load_file(w_path), strict=False) #Must be false 
+                 state_dict = load_file(w_path)
              else:
-                 self.model.load_state_dict(torch.load(os.path.join(path, "pytorch_model.bin"), map_location="cpu"), strict=True)
-        except:
-             # Fallback
-             self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
-             replace_layernorm_with_rmsnorm(self.model)
+                 state_dict = torch.load(os.path.join(path, "pytorch_model.bin"), map_location="cpu")
+            
+             # --- SETTING STRICT=FALSE AS REQUESTED ---
+             self.model.load_state_dict(state_dict, strict=False)
         
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -64,12 +73,24 @@ class EndpointHandler:
         inputs = data.pop("inputs", data)
         if isinstance(inputs, str): inputs = [inputs]
 
+        # Tokenize
         tokenized_inputs = self.tokenizer(inputs, return_tensors="pt", padding=True).to(self.device)
-        if "token_type_ids" in tokenized_inputs: del tokenized_inputs["token_type_ids"]
+        
+        # Remove harmful args
+        if "token_type_ids" in tokenized_inputs: 
+            del tokenized_inputs["token_type_ids"]
 
+        # Generate
         with torch.no_grad():
-            generated_ids = self.model.generate(**tokenized_inputs, max_new_tokens=128, num_beams=5, early_stopping=True)
+            generated_ids = self.model.generate(
+                **tokenized_inputs,
+                max_new_tokens=128,
+                num_beams=5,
+                early_stopping=True
+            )
 
+        # Decode & Fix
         decoded_outputs = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
         final_outputs = [fix_arabic_output(text) for text in decoded_outputs]
-        return [{"generated_text": text} for text in final_outputs]
+
+        return [{"generated_text": text} for text in final_outputs]