Update app.py

app.py

@@ -40,8 +40,9 @@ def process(model_A, model_B):
     input_text = "POST: "+ prompt+ "\n\n RESPONSE A: "+response_A+"\n\n RESPONSE B: "+response_B+"\n\n Which response is better? RESPONSE"
     x = tokenizer([input_text], return_tensors='pt').input_ids.to(device)
     y = model.generate(x, max_new_tokens=1)
-    prefered = tokenizer.batch_decode(y, skip_special_tokens=True)[0] 
-    return  prompt,df[df['model'].isin([model_A, model_B])], prefered
+    prefered = tokenizer.batch_decode(y, skip_special_tokens=True)[0]
+    result = model_A if prefered == 'A' else model_B
+    return  prompt,df[df['model'].isin([model_A, model_B])], result
 
 title = "Compare Instruction Models to see which one is more helpful"
 description = "This app compares the outputs of various open-source, instruction-trained models from a [dataset](https://huggingface.co/datasets/{OUTPUTS_DATASET}) of human demonstrations using a reward model trained on the [Stanford Human Preferences Dataset (SHP)](https://huggingface.co/datasets/stanfordnlp/SHP)"