Update app.py

app.py

@@ -29,9 +29,31 @@ model = PeftModel.from_pretrained(base_model, "unica/CLiMA")
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)
 
+prompt_instruction_drug_reviews = f"""Given a drug review enclosed in triple quotes and a pair of entities E1 corresponding to the drug name and E2 corresponding to the treated condition, classify the relation holding between E1 and E2.
+The relations are identified with 9 labels from 0 to 8. The meaning of the labels is the following:
+0 means that E1 causes E2
+1 means that E2 causes E1
+2 means that E1 enables E2
+3 means that E2 enables E1
+4 means that E1 prevents E2
+5 means that E2 prevents E1
+6 means that E1 hinders E2
+7 means that E2 hinders E1
+8 means that E1 and E2 are in a relation different than any of the previous ones.
+Given X the label that you predicted, for the output use the format LABEL: X
+"""
+
+
 # Format prompt
 def format_prompt(user_input, entity1, entity2):
-    return f"Identify causal relations in the following clinical narrative:\n\n{user_input}\n\nEntity 1: {entity1}\nEntity 2: {entity2}\n\nCausal relations:"
+#return f"Identify causal relations in the following clinical narrative:\n\n{user_input}\n\nEntity 1: {entity1}\nEntity 2: {entity2}\n\nCausal relations:"
+  text = user_input
+  prompt_text = f"Text:'''{text}'''"
+  e1 = entity1
+  e2 = entity2
+  prompt_entities = f"\nEntities: E1: '''{e1}''', E2: '''{e2}'''"
+  full_prompt = f"<USER> {prompt_instruction_drug_reviews} {prompt_text} {prompt_entities} <ASSISTANT>"
+  return full_prompt
 
 # Prediction function
 def generate_relations(text, entity1, entity2):