Update genra_incremental.py

genra_incremental.py

@@ -1,3 +1,4 @@
+import spaces
 import torch
 import numpy as np
 import pandas as pd
@@ -27,6 +28,7 @@ class LLMGenerator:
         self.tokenizer = tokenizer
         self.llm_name = llm_name
 
+    @spaces.GPU(duration=120)
     def generate_answer(self, texts, query, mode='validate'):
         template_texts =""
         for i, text in enumerate(texts):
@@ -62,6 +64,7 @@ class LLMGenerator:
 
 
 class QAIndexer:
+    @spaces.GPU(duration=60)
     def __init__(self, index_type, emb_model):
         self.document_embedder = SentenceTransformersDocumentEmbedder(model=emb_model)
         self.document_embedder.warm_up()
@@ -126,6 +129,7 @@ def rank_aggregation(aggregator, lists, k):
 
 
 class GenraPipeline:
+    @spaces.GPU(duration=60)
     def __init__(self, llm_name, emb_model, aggregator, contexts):
         self.qa_indexer = QAIndexer('in_memory', emb_model)
         self.qa_retriever = QARetriever(self.qa_indexer.document_store)