vhr1007 commited on
Commit
ce94de4
1 Parent(s): 7f062b7

method model call

Browse files
Files changed (1) hide show
  1. app.py +14 -4
app.py CHANGED
@@ -2,7 +2,8 @@ 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
@@ -30,7 +31,7 @@ logging.basicConfig(level=logging.INFO)
30
  huggingface_token = os.getenv('HUGGINGFACE_HUB_TOKEN')
31
  if huggingface_token:
32
  try:
33
- login(token=huggingface_token, add_to_git_credential=True, write_permission=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}")
@@ -45,10 +46,19 @@ 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
  cache_folder = os.path.join(hf_home_dir, "transformers_cache")
51
- encoder = SentenceTransformer('nomic-ai/nomic-embed-text-v1.5', cache_folder=cache_folder)
 
 
 
 
 
 
 
 
 
52
  logging.info("Successfully loaded the SentenceTransformer model.")
53
  except Exception as e:
54
  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 transformers import AutoTokenizer, AutoModel
6
+ from sentence_transformers import models, SentenceTransformer
7
  from services.qdrant_searcher import QdrantSearcher
8
  from services.openai_service import generate_rag_response
9
  from utils.auth import token_required
 
31
  huggingface_token = os.getenv('HUGGINGFACE_HUB_TOKEN')
32
  if huggingface_token:
33
  try:
34
+ login(token=huggingface_token, add_to_git_credential=True)
35
  logging.info("Successfully logged into Hugging Face Hub.")
36
  except Exception as e:
37
  logging.error(f"Failed to log into Hugging Face Hub: {e}")
 
46
  if not qdrant_url or not access_token:
47
  raise ValueError("Qdrant URL or Access Token is not set. Please set the QDRANT_URL and QDRANT_ACCESS_TOKEN environment variables.")
48
 
49
+ # Initialize the SentenceTransformer model with trust_remote_code using transformers
50
  try:
51
  cache_folder = os.path.join(hf_home_dir, "transformers_cache")
52
+
53
+ # Load the tokenizer and model with trust_remote_code=True
54
+ tokenizer = AutoTokenizer.from_pretrained('nomic-ai/nomic-embed-text-v1.5', trust_remote_code=True)
55
+ model = AutoModel.from_pretrained('nomic-ai/nomic-embed-text-v1.5', trust_remote_code=True)
56
+
57
+ # Wrap the model into a SentenceTransformer
58
+ word_embedding_model = models.Transformer(model_name_or_path='nomic-ai/nomic-embed-text-v1.5', model=model, tokenizer=tokenizer)
59
+ pooling_model = models.Pooling(word_embedding_model.get_word_embedding_dimension())
60
+ encoder = SentenceTransformer(modules=[word_embedding_model, pooling_model])
61
+
62
  logging.info("Successfully loaded the SentenceTransformer model.")
63
  except Exception as e:
64
  logging.error(f"Failed to load the SentenceTransformer model: {e}")