Update app.py

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer, BitsAndBytesConfig
 from peft import PeftModel, PeftConfig
 import torch
+import regex as re
 
 # Load PEFT adapter configuration
 peft_config = PeftConfig.from_pretrained("unica/CLiMA")
@@ -57,11 +58,38 @@ def format_prompt(user_input, entity1, entity2):
 
 # Prediction function
 def generate_relations(text, entity1, entity2):
+    answer_label_regex_pattern =  re.compile(r'LABEL:?\s?(\d+)')
+
+    
     prompt = format_prompt(text, entity1, entity2)
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(**inputs, max_new_tokens=256, do_sample=False)
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response[len(prompt):].strip()  # remove prompt from output if echoed
+    modelOut = response[len(prompt):].strip()  # remove prompt from output if echoed
+    answer_match = answer_label_regex_pattern.search(modelOut)
+    if answer_match:
+        if answer_match.group(1)=='0'
+            return f"""{entity1} causes  {entity2}"""
+        elif answer_match.group(1)=='1'
+            return f"""{entity2} causes  {entity1}"""
+        elif answer_match.group(1)=='2'
+            return f"""{entity1} enables  {entity2}"""
+        elif answer_match.group(1)=='3'
+            return f"""{entity2} enables  {entity1}"""
+        elif answer_match.group(1)=='4'
+            return f"""{entity1} prevents  {entity2}"""
+        elif answer_match.group(1)=='5'
+            return f"""{entity2} prevents  {entity1}"""
+        elif answer_match.group(1)=='6'
+            return f"""{entity1} hinders  {entity2}"""
+        elif answer_match.group(1)=='7'
+            return f"""{entity2} hinders  {entity1}"""
+        elif answer_match.group(1)=='8'
+            return f"""No causal relation between {entity1} and {entity2}"""
+    else:
+        return 'No causal relation could be extracted'
+            
+        
 
 # Gradio UI
 demo = gr.Interface(