update

pages/arena.py

@@ -22,7 +22,7 @@ def create_arena():
             submit_button = gr.Button("✨ Submit ✨")
     
             with gr.Row():
-                columns = [gr.Textbox(label=f"Column {i+1}", lines=10) for i in range(len(prompts))]
+                columns = [gr.Textbox(label=f"Prompt {i+1}", lines=10) for i in range(len(prompts))]
             
             content_list = [prompt + '\n{' + datapoint + '}\n\nsummary:' for prompt in prompts]
             model = get_model_batch_generation("Qwen/Qwen2-1.5B-Instruct")

utils/model.py

@@ -55,7 +55,7 @@ class Model(torch.nn.Module):
     def return_model(self):
         return self.pipeline
 
-    def gen(self, content_list, temp=0.1, max_length=500, streaming=False):
+    def gen(self, content_list, temp=0.001, max_length=500, streaming=False):
         # Convert list of texts to input IDs
         input_ids = self.tokenizer(content_list, return_tensors="pt", padding=True, truncation=True).input_ids.to(self.model.device)
 
@@ -74,7 +74,7 @@ class Model(torch.nn.Module):
                         return_dict_in_generate=True,
                         output_scores=True,
                         streamer=streamer):
-                    pass  # TextStreamer automatically handles the streaming, no need to manually handle the output
+                    yield output  # TextStreamer automatically handles the streaming, no need to manually handle the output
         else:
             outputs = self.model.generate(
                 input_ids,