upadte app.py

app.py

@@ -99,7 +99,7 @@ if torch.cuda.is_available():
 
     # question generation
     qg_tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-7b1")
-    qg_model = BloomForCausalLM.from_pretrained("bigscience/bloom-7b1", torch_dtype=torch.bfloat16)
+    qg_model = BloomForCausalLM.from_pretrained("bigscience/bloom-7b1", torch_dtype=torch.bfloat16).to('cuda')
     # qg_model = BloomForCausalLM.from_pretrained("bigscience/bloom-7b1", torch_dtype=torch.bfloat16).to(device)
     # qg_tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-7b1")
     # qg_model = BloomForCausalLM.from_pretrained("bigscience/bloom-7b1", torch_dtype=torch.bfloat16).to(device)
@@ -557,7 +557,7 @@ def get_google_search_results(api_key, search_engine_id, google_search, sort_dat
     return search_results
 
 
-@spaces.GPU
+# @spaces.GPU
 def averitec_search(claim, generate_question, speaker="they", check_date="2024-07-01", n_pages=1):  # n_pages=3
     # default config
     api_key = os.environ["GOOGLE_API_KEY"]
@@ -688,15 +688,15 @@ def generate_step2_reference_corpus(reference_file):
     return tokenized_corpus, prompt_corpus
 
 @spaces.GPU
-def decorate_with_questions(claim, retrieve_evidence, top_k=5):  # top_k=10, 100
+def decorate_with_questions(claim, retrieve_evidence, top_k=3):  # top_k=5, 10, 100
     #
     reference_file = "averitec/data/train.json"
     tokenized_corpus, prompt_corpus = generate_step2_reference_corpus(reference_file)
     prompt_bm25 = BM25Okapi(tokenized_corpus)
 
     # Define the bloom model:
-    accelerator = Accelerator()
-    accel_device = accelerator.device
+    # accelerator = Accelerator()
+    # accel_device = accelerator.device
     # device = "cuda:0" if torch.cuda.is_available() else "cpu"
     # tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-7b1")
     # model = BloomForCausalLM.from_pretrained(