wrap automodel for zerogpu

langchainapp.py

@@ -2,7 +2,7 @@
 import spaces
 from torch.nn import DataParallel  
 from torch import Tensor  
-from transformers import AutoTokenizer, AutoModel  
+# from transformers import AutoTokenizer, AutoModel  
 from huggingface_hub import InferenceClient  
 from openai import OpenAI  
 from langchain_community.embeddings import HuggingFaceInstructEmbeddings  
@@ -45,9 +45,14 @@ hf_token, yi_token = load_env_variables()
 
 @spaces.GPU  
 def load_model():  
+    # Import AutoModel within the function to avoid issues with pickling in multiprocessing  
+    from transformers import AutoModel, AutoTokenizer
     tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token, trust_remote_code=True)  
+  
     return AutoModel.from_pretrained(model_name, token=hf_token, trust_remote_code=True).to(device)  
   
+# Load model  
+nvidiamodel = load_model()   
 # Load model  
 nvidiamodel = load_model()
 # nvidiamodel.set_pooling_include_prompt(include_prompt=False)