Samuel Thomas commited on
Commit
ebdd994
·
1 Parent(s): 8693fa5
Files changed (1) hide show
  1. tools.py +2 -1
tools.py CHANGED
@@ -23,6 +23,7 @@ from transformers import BlipProcessor, BlipForQuestionAnswering, pipeline
23
  from io import BytesIO
24
  from sentence_transformers import SentenceTransformer
25
  from transformers import RagRetriever, RagTokenizer, RagSequenceForGeneration
 
26
 
27
 
28
  import os
@@ -105,7 +106,7 @@ llm_pipe = pipeline(
105
  )
106
 
107
  # Initialize RAG components
108
- tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-base", trust_remote_code=True)
109
  retriever = RagRetriever.from_pretrained(
110
  "facebook/rag-token-base",
111
  index_name="exact", # or "legacy" for legacy FAISS index
 
23
  from io import BytesIO
24
  from sentence_transformers import SentenceTransformer
25
  from transformers import RagRetriever, RagTokenizer, RagSequenceForGeneration
26
+ from transformers import AutoTokenizer, AutoModelWithLMHead
27
 
28
 
29
  import os
 
106
  )
107
 
108
  # Initialize RAG components
109
+ tokenizer = AutoTokenizer.from_pretrained("facebook/rag-token-base", trust_remote_code=True)
110
  retriever = RagRetriever.from_pretrained(
111
  "facebook/rag-token-base",
112
  index_name="exact", # or "legacy" for legacy FAISS index