Spaces:
Sleeping
Sleeping
Update src/rag_engine.py
Browse files- 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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 51 |
|
| 52 |
# CHECK 3: Standard Public Models
|
| 53 |
else:
|
| 54 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 55 |
|
| 56 |
except Exception as e:
|
| 57 |
logger.error(f"Failed to load embedding model '{model_name}': {e}")
|
| 58 |
-
return HuggingFaceEmbeddings(
|
|
|
|
|
|
|
|
|
|
| 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)
|