fix cpu

app.py

@@ -38,13 +38,23 @@ datasets = ["scifact"]
 current_dataset = "scifact"
 
 
-def pool(last_hidden_states, attention_mask):
+def pool(last_hidden_states, attention_mask, pool_type="last"):
     last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
-    sequence_lengths = attention_mask.sum(dim=1) - 1
-    batch_size = last_hidden.shape[0]
-    return last_hidden[torch.arange(batch_size, device=last_hidden.device), sequence_lengths]
 
-def create_batch_dict(tokenizer, input_texts, max_length=512):
+    if pool_type == "last":
+        left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
+        if left_padding:
+            emb = last_hidden[:, -1]
+        else:
+            sequence_lengths = attention_mask.sum(dim=1) - 1
+            batch_size = last_hidden.shape[0]
+            emb = last_hidden[torch.arange(batch_size, device=last_hidden.device), sequence_lengths]
+    else:
+        raise ValueError(f"pool_type {pool_type} not supported")
+
+    return emb
+
+def create_batch_dict(tokenizer, input_texts, always_add_eos="last", max_length=512):
     batch_dict = tokenizer(
         input_texts,
         max_length=max_length - 1,
@@ -53,7 +63,10 @@ def create_batch_dict(tokenizer, input_texts, max_length=512):
         padding=False,
         truncation=True
     )
-    batch_dict['input_ids'] = [input_ids + [tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
+
+    if always_add_eos == "last":
+        batch_dict['input_ids'] = [input_ids + [tokenizer.eos_token_id] for input_ids in batch_dict['input_ids']]
+
     return tokenizer.pad(
         batch_dict,
         padding=True,
@@ -62,18 +75,44 @@ def create_batch_dict(tokenizer, input_texts, max_length=512):
         return_tensors="pt",
     )
 
-def load_model():
-    global tokenizer, model, CUR_MODEL, BASE_MODEL
-    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-    tokenizer.pad_token_id = tokenizer.eos_token_id
-    tokenizer.pad_token = tokenizer.eos_token
-    tokenizer.padding_side = "right"
-    
-    # model = AutoModel.from_pretrained(CUR_MODEL, max_memory={"cpu": "12GiB"}, torch_dtype=torch.bfloat16, offload_state_dict=True)
-    base_model_instance = AutoModel.from_pretrained(BASE_MODEL).cpu()
-    model = PeftModel.from_pretrained(base_model_instance, CUR_MODEL)
-    model = model.merge_and_unload()
-    model.eval()
+class RepLlamaModel:
+    def __init__(self, model_name_or_path):
+        self.base_model = "meta-llama/Llama-2-7b-hf"
+        self.tokenizer = AutoTokenizer.from_pretrained(self.base_model)
+        self.tokenizer.model_max_length = 2048
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+        self.tokenizer.padding_side = "right"
+
+        self.model = self.get_model(model_name_or_path)
+        self.model.config.max_length = 2048
+
+    def get_model(self, peft_model_name):
+        base_model = AutoModel.from_pretrained(self.base_model)
+        model = PeftModel.from_pretrained(base_model, peft_model_name)
+        model = model.merge_and_unload()
+        model.eval()
+        return model
+
+    @spaces.GPU
+    def encode(self, texts, batch_size=32, **kwargs):
+        self.model = self.model.cuda()
+        all_embeddings = []
+        for i in range(0, len(texts), batch_size):
+            batch_texts = texts[i:i+batch_size]
+            
+            batch_dict = create_batch_dict(self.tokenizer, batch_texts, always_add_eos="last")
+            batch_dict = {key: value.cuda() for key, value in batch_dict.items()}
+
+            with torch.cuda.amp.autocast():
+                with torch.no_grad():
+                    outputs = self.model(**batch_dict)
+                    embeddings = pool(outputs.last_hidden_state, batch_dict['attention_mask'], 'last')
+                    embeddings = F.normalize(embeddings, p=2, dim=-1)
+                    all_embeddings.append(embeddings.cpu().numpy())
+
+        self.model = self.model.cpu()
+        return np.concatenate(all_embeddings, axis=0)
 
 
 def load_faiss_index(dataset_name):
@@ -128,31 +167,6 @@ def load_queries(dataset_name):
             qrels[dataset_name][qrel.query_id] = {}
         qrels[dataset_name][qrel.query_id][qrel.doc_id] = qrel.relevance
 
-@spaces.GPU
-def encode_queries(dataset_name, postfix):
-    global queries, tokenizer, model
-    input_texts = [f"query: {query.strip()} {postfix}".strip() for query in queries[dataset_name]]
-    
-    encoded_embeds = []
-    batch_size = 32
-    model = model.cuda()
-    
-    for start_idx in tqdm.tqdm(range(0, len(input_texts), batch_size), desc="Encoding queries"):
-        batch_input_texts = input_texts[start_idx: start_idx + batch_size]
-        
-        batch_dict = create_batch_dict(tokenizer, batch_input_texts)
-        batch_dict = {k: v.to(model.device) for k, v in batch_dict.items()}
-
-        with torch.cuda.amp.autocast():
-            with torch.no_grad():
-                outputs = model(**batch_dict)
-                embeds = pool(outputs.last_hidden_state, batch_dict['attention_mask'])
-                embeds = F.normalize(embeds, p=2, dim=-1)
-                encoded_embeds.append(embeds.float().cpu().numpy())
-
-    model = model.cpu()
-    return np.concatenate(encoded_embeds, axis=0)
-
 
 def evaluate(qrels, results, k_values):
     evaluator = pytrec_eval.RelevanceEvaluator(
@@ -168,15 +182,11 @@ def evaluate(qrels, results, k_values):
     return metrics
 
 def run_evaluation(dataset, postfix):
-    global current_dataset
-    
-    if dataset not in corpus_lookups or dataset not in queries:
-        load_corpus_lookups(dataset)
-        load_queries(dataset)
-    
+    global current_dataset, queries, model
     current_dataset = dataset
-    
-    q_reps = encode_queries(dataset, postfix)
+
+    input_texts = [f"query: {query.strip()} {postfix}".strip() for query in queries[current_dataset]]
+    q_reps = model.encode(input_texts)
     all_scores, psg_indices = search_queries(dataset, q_reps)
     
     results = {qid: dict(zip(doc_ids, map(float, scores))) 
@@ -189,16 +199,18 @@ def run_evaluation(dataset, postfix):
         "Recall@100": metrics["Recall@100"]
     }
 
-def gradio_interface(dataset, postfix):
-    if 'model' not in globals() or model is None:
-        load_model()
-        for dataset in datasets:
-            print(f"Loading dataset: {dataset}")
-            load_corpus_lookups(dataset)
-            load_queries(dataset)
 
+@spaces.GPU
+def gradio_interface(dataset, postfix):
     return run_evaluation(dataset, postfix)
 
+
+if model is None:
+    model = RepLlamaModel(model_name_or_path=CUR_MODEL)
+    load_corpus_lookups(current_dataset)
+    load_queries(current_dataset)   
+
+
 # Create Gradio interface
 iface = gr.Interface(
     fn=gradio_interface,