Update preference technique

app.py

@@ -7,11 +7,9 @@ from typing import Iterator
 import gradio as gr
 import spaces
 import torch  # noqa
-from transformers import (
-    AutoModelForCausalLM,  # noqa
-    AutoTokenizer,  # noqa
-    TextIteratorStreamer,  # noqa
-)
+from transformers import AutoModelForCausalLM  # noqa
+from transformers import AutoTokenizer  # noqa
+from transformers import TextIteratorStreamer  # noqa
 
 from chat_interface_preference import ChatInterface
 
@@ -70,7 +68,7 @@ def generate(
 
 chat_interface = ChatInterface(
     fn=generate,
-    prefence_techniques="dpo",
+    prefence_technique="dpo",
     min_turns=1,
     max_turns=10,
     repo_id="llm-human-feedback-collector-chat-interface-dpo",

chat_interface_preference.py

@@ -11,7 +11,7 @@ import json
 import random
 import re
 import uuid
-from typing import AsyncGenerator, Callable, List, Literal, Union, cast
+from typing import AsyncGenerator, Callable, Literal, Union, cast
 
 import anyio
 from gradio.blocks import Blocks
@@ -27,9 +27,8 @@ from gradio.components import (
     get_component_instance,
 )
 from gradio.events import Dependency, on
-from gradio.helpers import Error, Info
+from gradio.helpers import Error, Info, special_args
 from gradio.helpers import create_examples as Examples  # noqa: N812
-from gradio.helpers import special_args
 from gradio.layouts import Accordion, Group, Row
 from gradio.routes import Request
 from gradio.themes import ThemeClass as Theme
@@ -66,7 +65,7 @@ class ChatInterface(Blocks):
         self,
         fn: Callable,
         *,
-        prefence_techniques: str | List[str] | None = None,
+        prefence_technique: str = None,
         min_turns: int = 1,
         max_turns: int = 1,
         repo_id: None | str,
@@ -127,14 +126,9 @@ class ChatInterface(Blocks):
             raise ValueError("`max_turns` should be larger than `min_turns`")
         self.max_turns = max_turns
         self.min_turns = min_turns
-        if isinstance(prefence_techniques, str):
-            prefence_techniques = [prefence_techniques]
-        elif prefence_techniques is None:
-            prefence_techniques = ["sft"]
-        self.prefence_techniques = [technique.lower() for technique in prefence_techniques]
 
         optional_techniques = ["kto", "sft", "spin", "dpo", "simpo", "rlhf", "orpo"]
-        if any([technique for technique in self.prefence_techniques if technique not in optional_techniques]):
+        if prefence_technique not in optional_techniques:
             raise ValueError(f"Supported techniques are {optional_techniques}")
         submit_btn_one = "Generate"
         submit_btn_two = None
@@ -146,11 +140,12 @@ class ChatInterface(Blocks):
         stop_btn = "Stop"
         undo_btn = "↩️ Undo"
         clear_btn = "🗑️ Clear"
-        if "kto" in prefence_techniques:
+        n_generations = 1
+        if "kto" == prefence_technique:
             submit_btn_good = "The response 👍"
             submit_btn_bad = "The response 👎"
-        if any([technique for technique in ["dpo", "simpo", "rlhf", "orpo"] if technique in self.prefence_techniques]):
-            submit_btn_two = None
+        if prefence_technique in ["dpo", "simpo", "rlhf", "orpo"]:
+            n_generations = 2
             submit_btn_a = "A is better than B"
             submit_btn_b = "B is better than A"
             submit_btn_ab = "A and B are similar"
@@ -368,7 +363,7 @@ class ChatInterface(Blocks):
             self.saved_input = State()
             self.chatbot_state = State(self.chatbot.value) if self.chatbot.value else State([])
 
-            self._setup_events()
+            self._setup_events(n_generations)
             self._setup_api()
 
     def _set_conversation_id(self):
@@ -385,15 +380,15 @@ class ChatInterface(Blocks):
                 with self.data_file.open("a") as f:
                     f.write(json.dumps(feedback))
 
-    def _setup_events(self) -> None:
+    def _setup_events(self, n_generations) -> None:
         submit_fn_one = self._stream_fn if self.is_generator else self._submit_fn
-        submit_fn_one_partial = functools.partial(submit_fn_one, n_generations=2)
+        submit_fn_one_partial = functools.partial(submit_fn_one, n_generations=n_generations)
         submit_triggers_one = (
             [self.textbox.submit, self.submit_btn_one.click] if self.submit_btn_one else [self.textbox.submit]
         )
         submit_tuples = [(submit_fn_one_partial, submit_triggers_one)]
         if self.submit_btn_two:
-            submit_fn_two = functools.partial(submit_fn_one, n_generations=1)
+            submit_fn_two = functools.partial(submit_fn_one, n_generations=n_generations)
             submit_triggers_two = [self.submit_btn_two.click]
             submit_tuples.append((submit_fn_two, submit_triggers_two))
         for _fn, _triggers in submit_tuples: