Update models.py

models.py

@@ -13,7 +13,7 @@ You will be tasked to classify sentences as 'J' or 'V'
 
 Text: "{sentence}"
 
-Please classify this text as either 'J', 'W', or 'V'. Only output 'J', 'W', or 'V' with no additional explanation.<|endoftext|>
+Please classify this text as either 'J' or 'V'. Only output 'J' or 'V' with no additional explanation.<|endoftext|>
 <|assistant|>
 """
     return prompt
@@ -33,15 +33,42 @@ Please revise this text such that it maintains the criticism in the original tex
 """
     return prompt
 
+def mistral_score(sentence):
+    prompt = f"""<s>[INST]
+You are an assistant helping with paper reviews. 
+You will be tasked to classify sentences as 'J' or 'V'
+
+'J' is positive or 'J' is encouraging. 
+'J' has a neutral tone or 'J' is professional. 
+'V' is overly blunt or 'V' contains excessive negativity and no constructive feedback.
+'V' contains an accusatory tone or 'V' contains sweeping generalizations or 'V' contains personal attacks. 
+
+Text: "{sentence}"
+
+Please classify this text as either 'J' or 'V'. Only output 'J' or 'V' with no additional explanation. [/INST]"""
+    return prompt
+
+def mistral_revise(sentence):
+    prompt = f"""<s>[INST] 
+You are an assistant that helps users revise Paper Reviews.
+Paper reviews exist to provide authors of academic research papers constructive critism.
+
+This is text found in a review.
+This text was classified as 'toxic':
+
+Text: "{sentence}"
+
+Please revise this text such that it maintains the criticism in the original text and delivers it in a friendly but professional manner. Make minimal changes to the original text. [/INST] Revised Text: """
+    return prompt
+
 def query_model_score(sentence, api_key, model_id, prompt_fun):
     API_URL = f"https://api-inference.huggingface.co/models/{model_id}"
     headers = {"Authorization": f"Bearer {api_key}"}
     prompt = prompt_fun(sentence)
     def query(payload):
-        print(payload)
         response = requests.post(API_URL, headers=headers, json=payload)
         return response.json()
-    parameters = {"max_new_tokens" : 20, "temperature": 0.0, "return_full_text": False}
+    parameters = {"max_new_tokens" : 5, "temperature": 0.1, "return_full_text": False}
     options = {"wait_for_model": True}
     data = query({"inputs": f"{prompt}", "parameters": parameters, "options": options})
     score = data[0]['generated_text']
@@ -57,7 +84,7 @@ def query_model_revise(sentence, api_key, model_id, prompt_fun):
     def query(payload):
         response = requests.post(API_URL, headers=headers, json=payload)
         return response.json()
-    parameters = {"max_new_tokens" : 200, "temperature": 0.0, "return_full_text": False}
+    parameters = {"max_new_tokens" : 200, "temperature": 0.1, "return_full_text": False}
     options = {"wait_for_model": True}
     data = query({"inputs": f"{prompt}", "parameters": parameters, "options": options})
     revision = data[0]['generated_text']
@@ -75,6 +102,13 @@ def revise_review(review, api_key, model_id, highlight_color):
         "message": ""
     }
 
+    if 'zephyr' in model_id:
+        revision_prompt = zephyr_revise
+        score_prompt = zephyr_score
+    elif 'mistral' in model_id:
+        revision_prompt = mistral_revise
+        score_prompt = mistral_score
+
     try:
         review = review.replace('"', "'")
         sentences = parser.parse_sentences(review)
@@ -83,13 +117,13 @@ def revise_review(review, api_key, model_id, highlight_color):
         review_revision = ""
         for sentence in sentences:
             if len(sentence) > 20:
-                score = query_model_score(sentence, api_key, model_id, zephyr_score)
+                score = query_model_score(sentence, api_key, model_id, score_prompt)
                 if score == 0:
                     review_revision += " " + sentence
                 else:
                     review_score = 1
                     revision_count +=1
-                    revision = query_model_revise(sentence, api_key, model_id, zephyr_revise)
+                    revision = query_model_revise(sentence, api_key, model_id, revision_prompt)
                     revision = revision.strip().strip('"')
                     review_revision += f"<div style='background-color: {highlight_color}; display: inline;'>{revision}</div>"
             else: