import os
import shutil
import tempfile
import unittest

import torch
from modelscope import AutoTokenizer, Model

from swift import Swift
from swift.tuners.rome import RomeConfig


class TestRome(unittest.TestCase):

    def setUp(self):
        print(('Testing %s.%s' % (type(self).__name__, self._testMethodName)))
        self.tmp_dir = tempfile.TemporaryDirectory().name
        if not os.path.exists(self.tmp_dir):
            os.makedirs(self.tmp_dir)

    def tearDown(self):
        shutil.rmtree(self.tmp_dir)
        super().tearDown()

    @unittest.skip(
        'Rome test is skipped because the test image do not have flash-attn2')
    def test_rome(self):
        model = Model.from_pretrained(
            'modelscope/Llama-2-7b-ms',
            device_map='auto',
            trust_remote_code=True)
        tokenizer = AutoTokenizer.from_pretrained(
            'modelscope/Llama-2-7b-ms', trust_remote_code=True)
        request = [{
            'prompt': '{} was the founder of',
            'subject': 'Steve Jobs',
            'target': 'Microsoft',
        }]
        config = RomeConfig(
            model_type='llama-7b',
            knowledge=request,
            tokenizer=tokenizer,
        )

        model = Swift.prepare_model(model, config)
        prompt = 'Steve Jobs was the founder of'
        inp_tok = tokenizer(
            prompt, return_token_type_ids=False, return_tensors='pt')
        for key, value in inp_tok.items():
            inp_tok[key] = value.to('cuda')
        with torch.no_grad():
            generated_ids = model.generate(
                **inp_tok,
                temperature=0.1,
                top_k=50,
                max_length=128,
                do_sample=True)

        responses = tokenizer.batch_decode(
            generated_ids[:, inp_tok['input_ids'].size(1):],
            skip_special_tokens=True,
            clean_up_tokenization_spaces=True)
        self.assertTrue('Microsoft' in responses[0])
