add correlation score

app.py

@@ -1,23 +1,48 @@
 import torch
 import transformers
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForSequenceClassification
 import gradio as gr
 import os
 
 model_name = 'eliolio/bart-finetuned-yelpreviews'
+bert_model_name = 'eliolio/bert-correlation-yelpreviews'
 
 access_token = os.environ.get('private_token')
 
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name, use_auth_token=access_token)
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=access_token)
+model = AutoModelForSeq2SeqLM.from_pretrained(
+    model_name, use_auth_token=access_token
+)
+tokenizer = AutoTokenizer.from_pretrained(
+    model_name, use_auth_token=access_token
+)
+
+bert_tokenizer = AutoTokenizer.from_pretrained(
+    bert_model_name, use_auth_token=access_token
+)
+bert_model = AutoModelForSequenceClassification.from_pretrained(
+    bert_model_name, use_auth_token=access_token
+)
+
+
+def correlation_score(table, review):
+    # Compute the correlation score
+    args = ((table, review))
+    inputs = bert_tokenizer(*args, padding=True, max_length=128, truncation=True, return_tensors="pt")
+    logits = model(**inputs).logits
+    probs = logits.softmax(dim=-1)
+    return {
+        "correlation": probs[:, 1].item()
+    }
 
 def create_prompt(stars, useful, funny, cool):
     return f"Generate review: stars: {stars}, useful: {useful}, funny: {funny}, cool: {cool}"
 
+
 def postprocess(review):
     dot = review.rfind('.')
     return review[:dot+1]
 
+
 def generate_reviews(stars, useful, funny, cool):
     text = create_prompt(stars, useful, funny, cool)
     inputs = tokenizer(text, return_tensors='pt')
@@ -30,10 +55,15 @@ def generate_reviews(stars, useful, funny, cool):
         top_p=0.9
     )
     reviews = []
+    scores = []
     for review in out:
         reviews.append(postprocess(tokenizer.decode(review, skip_special_tokens=True)))
+        scores.append(
+            correlation_score(text[17:], review)
+        )
+
+    return reviews[0], reviews[1], reviews[2], scores[0], scores[1], scores[2]
 
-    return reviews[0], reviews[1], reviews[2]
 
 css = """
     #ctr {text-align: center;}
@@ -65,12 +95,16 @@ demo = gr.Blocks(css=css)
 with demo:
     with gr.Row():
         gr.Markdown(md_text)
-    
+
     with gr.Row():
-        stars = gr.inputs.Slider(minimum=0, maximum=5, step=1, default=0, label="stars")
-        useful = gr.inputs.Slider(minimum=0, maximum=5, step=1, default=0, label="useful")
-        funny = gr.inputs.Slider(minimum=0, maximum=5, step=1, default=0, label="funny")
-        cool = gr.inputs.Slider(minimum=0, maximum=5, step=1, default=0, label="cool")
+        stars = gr.inputs.Slider(minimum=0, maximum=5,
+                                 step=1, default=0, label="stars")
+        useful = gr.inputs.Slider(
+            minimum=0, maximum=5, step=1, default=0, label="useful")
+        funny = gr.inputs.Slider(minimum=0, maximum=5,
+                                 step=1, default=0, label="funny")
+        cool = gr.inputs.Slider(minimum=0, maximum=5,
+                                step=1, default=0, label="cool")
     with gr.Row():
         button = gr.Button("Generate reviews !", elem_id='btn')
 
@@ -79,13 +113,18 @@ with demo:
         output2 = gr.Textbox(label="Review #2")
         output3 = gr.Textbox(label="Review #3")
 
+    with gr.Row():
+        score1 = gr.Label(label="Correlation score #1")
+        score2 = gr.Label(label="Correlation score #2")
+        score3 = gr.Label(label="Correlation score #3")
+
     with gr.Row():
         gr.Markdown(resources)
 
     button.click(
         fn=generate_reviews,
         inputs=[stars, useful, funny, cool],
-        outputs=[output1, output2, output3]
+        outputs=[output1, output2, output3, score1, score2, score3]
     )
 
-demo.launch()
+demo.launch()