NavyDevilDoc commited on
Commit
10a9f86
·
verified ·
1 Parent(s): 15383a5

Update src/rag_engine.py

Browse files
Files changed (1) hide show
  1. src/rag_engine.py +15 -3
src/rag_engine.py CHANGED
@@ -47,15 +47,27 @@ def get_embedding_func(model_name: str = "sentence-transformers/all-MiniLM-L6-v2
47
  allow_patterns=f"{folder_name}/*"
48
  )
49
  local_model_path = os.path.join(storage_path, folder_name)
50
- return HuggingFaceEmbeddings(model_name=local_model_path)
 
 
 
 
 
51
 
52
  # CHECK 3: Standard Public Models
53
  else:
54
- return HuggingFaceEmbeddings(model_name=model_name)
 
 
 
 
55
 
56
  except Exception as e:
57
  logger.error(f"Failed to load embedding model '{model_name}': {e}")
58
- return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
 
 
59
 
60
  def save_uploaded_file(uploaded_file, username: str) -> str:
61
  user_dir = os.path.join(UPLOAD_DIR, username)
 
47
  allow_patterns=f"{folder_name}/*"
48
  )
49
  local_model_path = os.path.join(storage_path, folder_name)
50
+
51
+ # FIX: Explicitly set device to CPU to avoid meta-tensor errors
52
+ return HuggingFaceEmbeddings(
53
+ model_name=local_model_path,
54
+ model_kwargs={'device': 'cpu'}
55
+ )
56
 
57
  # CHECK 3: Standard Public Models
58
  else:
59
+ # FIX: Explicitly set device to CPU here as well
60
+ return HuggingFaceEmbeddings(
61
+ model_name=model_name,
62
+ model_kwargs={'device': 'cpu'}
63
+ )
64
 
65
  except Exception as e:
66
  logger.error(f"Failed to load embedding model '{model_name}': {e}")
67
+ return HuggingFaceEmbeddings(
68
+ model_name="sentence-transformers/all-MiniLM-L6-v2",
69
+ model_kwargs={'device': 'cpu'}
70
+ )
71
 
72
  def save_uploaded_file(uploaded_file, username: str) -> str:
73
  user_dir = os.path.join(UPLOAD_DIR, username)