"""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()