sanjeevbora commited on
Commit
7bc17d6
·
verified ·
1 Parent(s): 90dbd88

updated with gpu

Browse files
Files changed (1) hide show
  1. app.py +2 -0
app.py CHANGED
@@ -7,6 +7,7 @@ from transformers import AutoConfig, AutoTokenizer, pipeline, AutoModelForCausal
7
  import torch
8
  import re
9
  import transformers
 
10
 
11
  # Initialize embeddings and ChromaDB
12
  model_name = "sentence-transformers/all-mpnet-base-v2"
@@ -62,6 +63,7 @@ books_db_client_retriever = RetrievalQA.from_chain_type(
62
  )
63
 
64
  # Function to retrieve answer using the RAG system
 
65
  def test_rag(query):
66
  books_retriever = books_db_client_retriever.run(query)
67
 
 
7
  import torch
8
  import re
9
  import transformers
10
+ import spaces
11
 
12
  # Initialize embeddings and ChromaDB
13
  model_name = "sentence-transformers/all-mpnet-base-v2"
 
63
  )
64
 
65
  # Function to retrieve answer using the RAG system
66
+ @spaces.GPU(duration=120)
67
  def test_rag(query):
68
  books_retriever = books_db_client_retriever.run(query)
69