init

app/embedding_setup.py

@@ -32,12 +32,13 @@ embedding_sim = HuggingFaceBgeEmbeddings(
 db = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embedding_int)
 retriever = db.as_retriever(search_kwargs={"k": TOP_K})
 
-
+lorax_client = pb.deployments.client("llama-3-8b-instruct") # Insert deployment name here
 
 lora_weights_rec = REC_LORA_MODEL
 lora_weights_exp = EXP_LORA_MODEL
 hf_auth  = os.environ.get("hf_token")
 
+
 tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL, token=hf_auth)
 
 
@@ -46,28 +47,6 @@ second_token = 'Second'
 # 获取token的ID
 first_id = tokenizer.convert_tokens_to_ids(first_token)
 second_id = tokenizer.convert_tokens_to_ids(second_token)
-model = AutoModelForCausalLM.from_pretrained(
-            LLM_MODEL,
-            load_in_4bit=True,
-            torch_dtype=torch.float16,
-            token=hf_auth,
-        )
-
-rec_adapter = PeftModel.from_pretrained(
-            model,
-            lora_weights_rec
-        )
-      
-
-tokenizer.padding_side = "left"
-    # unwind broken decapoda-research config
-#model.half()  # seems to fix bugs for some users.
-rec_adapter.eval()
-
-rec_adapter.config.pad_token_id = tokenizer.pad_token_id = 0  # unk
-rec_adapter.config.bos_token_id = 1
-rec_adapter.config.eos_token_id = 2
-
 
 
 def generate_prompt(target_occupation, skill_gap, courses):
@@ -100,32 +79,9 @@ def evaluate(
         **kwargs,
     ):
 
-        inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(device)
-        generation_config = GenerationConfig(
-            temperature=temperature,
-            top_p=top_p,
-            top_k=top_k,
-            num_beams=num_beams,
-            **kwargs,
-        )
-        with torch.no_grad():
-            rec_adapter.to(device)
-            generation_output = rec_adapter.generate(
-                **inputs,
-                generation_config=generation_config,
-                return_dict_in_generate=True,
-                output_scores=True,
-                max_new_tokens=max_new_tokens,
-                # batch_size=batch_size,
-                eos_token_id=tokenizer.eos_token_id,
-                pad_token_id=tokenizer.eos_token_id,
-            )
-        scores = generation_output.scores[0].softmax(dim=-1)
-        logits = torch.tensor(scores[:,[first_id, second_id]], dtype=torch.float32).softmax(dim=-1)
-        s = generation_output.sequences
-        output = tokenizer.batch_decode(s, skip_special_tokens=True)
-        output = [_.split('Response:\n')[-1] for _ in output]
-        return output, logits.tolist()
+        resp = lorax_client.generate(prompt,adapter_id=REC_LORA_MODEL, adapter_source='hub', max_new_tokens=max_new_tokens)
+        
+        return resp
 
 def compare_docs_with_context(doc_a, doc_b, target_occupation_name, target_occupation_dsp,skill_gap):
     
@@ -134,13 +90,14 @@ def compare_docs_with_context(doc_a, doc_b, target_occupation_name, target_occup
     target_occupation = f"name: {target_occupation_name} description: {target_occupation_dsp[:1500]}"
     skill_gap = skill_gap
     prompt = generate_prompt(target_occupation, skill_gap, courses)
-    prompt = [prompt]
-    output, logit = evaluate(prompt)
+    prompt = prompt
+    output = evaluate(prompt)
     # Compare based on the response: [A] means doc_a > doc_b, [B] means doc_a < doc_b
-    print(output, logit)
-    if logit[0][0] > logit[0][1]:
+    print(output)
+    result_token_id = output.details.token[0].id
+    if result_token_id == first_id:
         return 1  # doc_a should come before doc_b
-    elif logit[0][0] < logit[0][1]:
+    elif result_token_id == second_id:
         return -1  # doc_a should come after doc_b
     else:
         return 0  # Consider them equal if the response is unclear
@@ -148,7 +105,7 @@ def compare_docs_with_context(doc_a, doc_b, target_occupation_name, target_occup
 
 #-----------------------------------------explanation-------------------------------------
 
-lorax_client = pb.deployments.client("llama-3-8b-instruct") # Insert deployment name here
+
 def generate_prompt_exp(input_text):
     return f"""
 ### Instruction: