from pre trained

UModel.py

@@ -153,6 +153,35 @@ class UModel(nn.Module):
 
         predstrs = [uyghur_latin.decode(pred) for pred in preds]
         return predstrs
+    
+    @classmethod
+    def from_pretrained(cls, repo_id: str, filename: str = "best_model.pth", device: str = "cpu", **kwargs):
+        """
+        Load a pretrained UModel from a Hugging Face repository.
+
+        Args:
+            repo_id (str): The Hugging Face repository ID.
+            filename (str): The name of the checkpoint file in the repo.
+            device (str): The device to load the model onto ("cpu" or "cuda").
+            **kwargs: Additional arguments to pass to the UModel constructor.
+
+        Returns:
+            UModel: An instance of UModel with pretrained weights loaded.
+        """
+        # Download the checkpoint from the Hugging Face Hub
+        checkpoint_path = hf_hub_download(repo_id=repo_id, filename=filename)
+        
+        # Initialize the model
+        model = cls(**kwargs)
+        
+        # Load state dict
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+        model.load_state_dict(checkpoint["st_dict"])
+        
+        # Set model to evaluation mode
+        model.eval()
+        print(f"Loaded model from {repo_id}/{filename} onto {device}")
+        return model
 
 
 if __name__ == "__main__":