File size: 1,641 Bytes
be32ef5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 | """GPU-host tests; local execution requires torch, not installed by this test."""
import sys
from pathlib import Path
import unittest
sys.path.insert(0,str(Path(__file__).resolve().parents[1]/'src'))
import torch
from visual_preference import sequence_logps,anchored_loss
class Preference(unittest.TestCase):
def test_completion_mask_and_shift(self):
logits=torch.zeros(1,4,3,requires_grad=True)
labels=torch.tensor([[-100,-100,1,2]])
lp=sequence_logps(logits,labels)
self.assertAlmostEqual(float(lp[0]),-1.098612288,places=6)
(-lp.mean()).backward()
self.assertEqual(float(logits.grad[0,0].abs().sum()),0.)
self.assertGreater(float(logits.grad[0,1].abs().sum()),0.)
self.assertGreater(float(logits.grad[0,2].abs().sum()),0.)
self.assertEqual(float(logits.grad[0,3].abs().sum()),0.)
def test_preference_gradient_direction(self):
good=torch.tensor([-.7],requires_grad=True);bad=torch.tensor([-1.],requires_grad=True)
loss,_,_=anchored_loss(good,bad,torch.ones(1));loss.backward()
self.assertLess(float(good.grad[0]),0.)
self.assertGreater(float(bad.grad[0]),0.)
def test_ineligible_pairs_have_only_sft_gradient(self):
good=torch.tensor([-1.],requires_grad=True);bad=torch.tensor([-1.],requires_grad=True)
loss,_,_=anchored_loss(good,bad,torch.zeros(1));loss.backward()
self.assertAlmostEqual(float(loss),1.)
self.assertEqual(float(bad.grad[0]),0.)
def test_length_normalization(self):
logits=torch.zeros(2,4,3)
labels=torch.tensor([[-100,1,2,-100],[-100,1,2,1]])
self.assertTrue(torch.allclose(sequence_logps(logits,labels),torch.full((2,),-1.098612288)))
if __name__=='__main__':unittest.main()
|