vhr1007 commited on
Commit
b619001
·
1 Parent(s): 4803e02
Files changed (1) hide show
  1. app.py +14 -9
app.py CHANGED
@@ -2,7 +2,7 @@ from huggingface_hub import login
2
  from fastapi import FastAPI, Depends, HTTPException
3
  import logging
4
  from pydantic import BaseModel
5
- from sentence_transformers import SentenceTransformer
6
  from services.qdrant_searcher import QdrantSearcher
7
  from services.openai_service import generate_rag_response
8
  from utils.auth import token_required
@@ -15,13 +15,13 @@ load_dotenv()
15
  # Initialize FastAPI application
16
  app = FastAPI()
17
 
 
18
  os.environ["HF_HOME"] = "/tmp/huggingface_cache"
19
 
20
-
21
  # Ensure the cache directory exists
22
- cache_dir = os.environ["HF_HOME"]
23
- if not os.path.exists(cache_dir):
24
- os.makedirs(cache_dir)
25
 
26
  # Setup logging
27
  logging.basicConfig(level=logging.INFO)
@@ -29,11 +29,15 @@ logging.basicConfig(level=logging.INFO)
29
  # Load Hugging Face token from environment variable
30
  huggingface_token = os.getenv('HUGGINGFACE_HUB_TOKEN')
31
  if huggingface_token:
32
- login(token=huggingface_token, add_to_git_credential=True, write_permission=True)
 
 
 
 
 
33
  else:
34
  raise ValueError("Hugging Face token is not set. Please set the HUGGINGFACE_HUB_TOKEN environment variable.")
35
 
36
-
37
  # Initialize the Qdrant searcher
38
  qdrant_url = os.getenv('QDRANT_URL')
39
  access_token = os.getenv('QDRANT_ACCESS_TOKEN')
@@ -41,9 +45,10 @@ access_token = os.getenv('QDRANT_ACCESS_TOKEN')
41
  if not qdrant_url or not access_token:
42
  raise ValueError("Qdrant URL or Access Token is not set. Please set the QDRANT_URL and QDRANT_ACCESS_TOKEN environment variables.")
43
 
44
- # Initialize the SentenceTransformer model
45
  try:
46
- encoder = SentenceTransformer('nomic-ai/nomic-embed-text-v1.5', trust_remote_code=True)
 
47
  logging.info("Successfully loaded the SentenceTransformer model.")
48
  except Exception as e:
49
  logging.error(f"Failed to load the SentenceTransformer model: {e}")
 
2
  from fastapi import FastAPI, Depends, HTTPException
3
  import logging
4
  from pydantic import BaseModel
5
+ from sentence_transformers import SentenceTransformer, logging as st_logging
6
  from services.qdrant_searcher import QdrantSearcher
7
  from services.openai_service import generate_rag_response
8
  from utils.auth import token_required
 
15
  # Initialize FastAPI application
16
  app = FastAPI()
17
 
18
+ # Set the cache directory for Hugging Face
19
  os.environ["HF_HOME"] = "/tmp/huggingface_cache"
20
 
 
21
  # Ensure the cache directory exists
22
+ hf_home_dir = os.environ["HF_HOME"]
23
+ if not os.path.exists(hf_home_dir):
24
+ os.makedirs(hf_home_dir)
25
 
26
  # Setup logging
27
  logging.basicConfig(level=logging.INFO)
 
29
  # Load Hugging Face token from environment variable
30
  huggingface_token = os.getenv('HUGGINGFACE_HUB_TOKEN')
31
  if huggingface_token:
32
+ try:
33
+ login(token=huggingface_token, add_to_git_credential=True)
34
+ logging.info("Successfully logged into Hugging Face Hub.")
35
+ except Exception as e:
36
+ logging.error(f"Failed to log into Hugging Face Hub: {e}")
37
+ raise HTTPException(status_code=500, detail="Failed to log into Hugging Face Hub.")
38
  else:
39
  raise ValueError("Hugging Face token is not set. Please set the HUGGINGFACE_HUB_TOKEN environment variable.")
40
 
 
41
  # Initialize the Qdrant searcher
42
  qdrant_url = os.getenv('QDRANT_URL')
43
  access_token = os.getenv('QDRANT_ACCESS_TOKEN')
 
45
  if not qdrant_url or not access_token:
46
  raise ValueError("Qdrant URL or Access Token is not set. Please set the QDRANT_URL and QDRANT_ACCESS_TOKEN environment variables.")
47
 
48
+ # Initialize the SentenceTransformer model with the cache directory managed by HF_HOME
49
  try:
50
+ st_logging.set_verbosity_info()
51
+ encoder = SentenceTransformer('nomic-ai/nomic-embed-text-v1.5')
52
  logging.info("Successfully loaded the SentenceTransformer model.")
53
  except Exception as e:
54
  logging.error(f"Failed to load the SentenceTransformer model: {e}")