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  
@@ -43,18 +43,20 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 ### Utils  
 hf_token, yi_token = load_env_variables()  
 
+tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token, trust_remote_code=True)  
+# Lazy load model  
+model = None  
+  
 @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)  
+    global model  
+    if model is None:  
+        from transformers import AutoModel  
+        model = AutoModel.from_pretrained(model_name, token=hf_token, trust_remote_code=True).to(device)  
+    return model  
   
 # Load model  
 nvidiamodel = load_model()   
-# Load model  
-nvidiamodel = load_model()
 # nvidiamodel.set_pooling_include_prompt(include_prompt=False) 
 
 def clear_cuda_cache():  
@@ -66,7 +68,8 @@ chroma_client = chromadb.Client(Settings())
   
 # Create a collection  
 chroma_collection = chroma_client.create_collection("all-my-documents")  
-  
+
+@spaces.GPU  
 class MyEmbeddingFunction(EmbeddingFunction):  
     def __init__(self, model_name: str, token: str, intention_client):  
         self.model_name = model_name  
@@ -156,6 +159,7 @@ def add_documents_to_chroma(documents: list, embedding_function: MyEmbeddingFunc
             )  
   
 def query_chroma(query_text: str, embedding_function: MyEmbeddingFunction):  
+    model = load_model()  
     query_embeddings, query_metadata = embedding_function.compute_embeddings(query_text)  
     result_docs = chroma_collection.query(  
         query_texts=[query_text],  
@@ -208,6 +212,7 @@ def upload_documents(files):
     return "Documents uploaded and processed successfully!"  
   
 def query_documents(query):  
+    model = load_model()  
     results = query_chroma(query)  
     return "\n\n".join([result.content for result in results])