get_ipython().run_line_magic('load_ext', 'autoreload')
get_ipython().run_line_magic('autoreload', '2')


import dspy
from dspy.evaluate import Evaluate
from dspy.datasets.hotpotqa import HotPotQA
from dspy.teleprompt import BootstrapFewShotWithRandomSearch, BootstrapFinetune



ports = [7140, 7141, 7142, 7143, 7144, 7145]
llamaChat = dspy.HFClientTGI(model="meta-llama/Llama-2-13b-chat-hf", port=ports, max_tokens=150)
colbertv2 = dspy.ColBERTv2(url='http://20.102.90.50:2017/wiki17_abstracts')

dspy.settings.configure(rm=colbertv2, lm=llamaChat)



dataset = HotPotQA(train_seed=1, train_size=200, eval_seed=2023, dev_size=1000, test_size=0)
trainset = [x.with_inputs('question') for x in dataset.train]
devset = [x.with_inputs('question') for x in dataset.dev]
testset = [x.with_inputs('question') for x in dataset.test]

len(trainset), len(devset), len(testset)


trainset[0]



from dsp.utils.utils import deduplicate

class BasicMH(dspy.Module):
    def __init__(self, passages_per_hop=3):
        super().__init__()

        self.retrieve = dspy.Retrieve(k=passages_per_hop)
        self.generate_query = [dspy.ChainOfThought("context, question -> search_query") for _ in range(2)]
        self.generate_answer = dspy.ChainOfThought("context, question -> answer")
    
    def forward(self, question):
        context = []
        
        for hop in range(2):
            search_query = self.generate_query[hop](context=context, question=question).search_query
            passages = self.retrieve(search_query).passages
            context = deduplicate(context + passages)

        return self.generate_answer(context=context, question=question).copy(context=context)



RECOMPILE_INTO_LLAMA_FROM_SCRATCH = False
NUM_THREADS = 24

metric_EM = dspy.evaluate.answer_exact_match


if RECOMPILE_INTO_LLAMA_FROM_SCRATCH:
    tp = BootstrapFewShotWithRandomSearch(metric=metric_EM, max_bootstrapped_demos=2, num_threads=NUM_THREADS)
    basicmh_bs = tp.compile(BasicMH(), trainset=trainset[:50], valset=trainset[50:200])

    ensemble = [prog for *_, prog in basicmh_bs.candidate_programs[:4]]

    for idx, prog in enumerate(ensemble):
        pass


if not RECOMPILE_INTO_LLAMA_FROM_SCRATCH:
    ensemble = []

    for idx in range(4):
        prog = BasicMH()
        prog.load(f'multihop_llama213b_{idx}.json')
        ensemble.append(prog)


llama_program = ensemble[0]

evaluate_hotpot = Evaluate(devset=devset[:1000], metric=metric_EM, num_threads=NUM_THREADS, display_progress=True, display_table=0)
evaluate_hotpot(llama_program)


llama_program(question="How many storeys are in the castle that David Gregory inherited?")

llamaChat.inspect_history(n=3)



unlabeled_train = HotPotQA(train_seed=1, train_size=3000, eval_seed=2023, dev_size=0, test_size=0).train
unlabeled_train = [
dspy.Example(question=x.question)