Update vlog4chat.py

vlog4chat.py

@@ -28,15 +28,20 @@ from pydantic.v1 import Extra, Field, root_validator
 from typing import Any, Dict, List, Mapping, Optional
 from langchain.memory import ConversationBufferMemory
 from langchain import LLMChain, PromptTemplate
-#from paddleocr import PaddleOCR, draw_ocr
 
-#sys.path.append('/root/autodl-tmp/recognize-anything')
+import torch
 
-#from ram.models import ram
-#from ram.models import tag2text
-#from ram import inference_ram as inference
-#from ram import inference_tag2text as inference
-#from ram import get_transform
+def map_to_cpu(obj):
+    if isinstance(obj, torch.Tensor):
+        return obj.to(torch.device('cpu'))
+    elif isinstance(obj, dict):
+        return {k: map_to_cpu(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [map_to_cpu(i) for i in obj]
+    elif isinstance(obj, tuple):
+        return tuple(map_to_cpu(i) for i in obj)
+    else:
+        return obj
 
 warnings.filterwarnings("ignore", category=UserWarning)
 B_INST, E_INST = "[INST]", "[/INST]"
@@ -166,8 +171,8 @@ class Vlogger4chat :
             #chunks = text_splitter.split_documents(raw_documents)
             #self.vector_storage = FAISS.from_documents(chunks, self.my_embedding)    
             with open('./BV11H4y1F7uH.pkl', 'rb') as f:
-                self.vector_storage =pickle.load(f)
-             
+                vector_storage =pickle.load(f)
+            self.vector_storage = map_to_cpu(vector_storage)
             self.chain = ConversationalRetrievalChain.from_llm(self.llm, self.vector_storage.as_retriever(), return_source_documents=True)
             return True
         return False