New model training interface

app.py

@@ -6,7 +6,7 @@ import torch
 import re
 import tempfile
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from repeng import ControlVector, ControlModel
+from repeng import ControlVector, ControlModel, DatasetEntry
 import gradio as gr
 
 # Initialize model and tokenizer
@@ -35,6 +35,7 @@ if cuda:
     
 model = ControlModel(model, list(range(-5, -18, -1)))
 
+# Generation settings
 # Generation settings
 default_generation_settings = {
     "pad_token_id": tokenizer.eos_token_id,
@@ -382,16 +383,16 @@ def train_model_persona(positive_text, negative_text):
         positive_list,
         negative_list,
         output_suffixes)
-#    model.reset()
-#    output_model = ControlVector.train(model, tokenizer, dataset
+    model.reset()
+    output_model = ControlVector.train(model, tokenizer, dataset)
     # Write file to temporary directory returning the path to Gradio for download
-    filename = re.sub(r'[<>:"/\\|?*]', '', positive_text) + '_'
+    filename = re.sub(r'[ <>:"/\\|?*]', '', positive_text) + '_'
     temp_file = tempfile.NamedTemporaryFile(
         prefix=filename,
         suffix=".gguf",
         delete= False
         )
-#    ControlVector.export_gguf(output_model, temp_file.name)
+    ControlVector.export_gguf(output_model, temp_file.name)
     temp_file.close()
     return temp_file.name
 
@@ -408,7 +409,7 @@ def train_model_facts(positive_text, negative_text):
     )
 
     output_model = ControlVector.train(model, tokenizer, dataset)
-    filename = re.sub(r'[<>:"/\\|?*]', '', positive_text) + '_'
+    filename = re.sub(r'[ <>:"/\\|?*]', '', positive_text) + '_'
     temp_file = tempfile.NamedTemporaryFile(
         prefix=filename,
         suffix=".gguf",
@@ -787,10 +788,10 @@ with gr.Blocks(
                 gr.Markdown("Fill in the blank with a persona and its opposite within, \"Pretend to be a (persona) making statements about the world.\"")
                 facts_input_positive = gr.Text(
                     label="Positive",
-                    placeholder="time traveller from the future")
+                    placeholder="time traveler from the future")
                 facts_input_negative = gr.Text(
                     label="Negative",
-                    placeholder="time travaller from the past")
+                    placeholder="time travaler from the past")
                 button_facts = gr.Button(
                     value="Generate fact control model"
                 )