How can we adopt this to use online instead of offline infrence? 
It should be an option and not a replacement, so it should be a simple change in the generate function.

class LLMCheck:
    def __init__(self, model_id, tensor_parallel_size=1, max_tokens=512, cache_dir=None, enable_prefix_caching=False, max_model_len=None):
        from vllm import LLM, SamplingParams
        from vllm.sampling_params import GuidedDecodingParams


        logger.info("Reminder: Please set the CUDA device before initializing the LLMCheck object.")

        self.model_id = model_id

        self.tensor_parallel_size = tensor_parallel_size
        self.max_tokens = max_tokens
        self.max_model_len = 32768 if max_model_len is None else max_model_len # max input length (prompt + doc)
        self.default_chunk_size = self.max_model_len - 300 # reserve some space (hard coded) for the claim to be checked
        self.cache_dir = cache_dir

        self.user_prompt = USER_PROMPT
        self.system_prompt = SYSTEM_PROMPT
        self.enable_prefix_caching = enable_prefix_caching

        # Check if CUDA is available and get compute capability
        if torch.cuda.is_available():
            compute_capability = torch.cuda.get_device_capability()
            if compute_capability[0] >= 8:
                self.dtype = torch.bfloat16
                logger.info("Using bfloat16 for LLM initialization.")
            else:
                self.dtype = torch.float16
                logger.info(f"GPU compute capability {compute_capability} < 8.0. Using float16 for LLM initialization.")
        else:
            if torch.cpu.is_available() and hasattr(torch.cpu, 'is_bf16_supported') and torch.cpu.is_bf16_supported():
                self.dtype = torch.bfloat16
                logger.info("CUDA not available. Using bfloat16 on CPU.")
            else:
                self.dtype = torch.float32
                logger.info("CUDA not available and CPU doesn't support bfloat16. Using float32 for LLM initialization.")
        
        self.llm = LLM(
            model=self.model_id, 
            dtype=self.dtype, 
            download_dir=self.cache_dir,
            trust_remote_code=True if self.model_id == 'bespokelabs/Bespoke-MiniCheck-7B' else False,
            tensor_parallel_size=self.tensor_parallel_size,
            seed=2024,
            max_model_len=self.max_model_len,   # need to be adjusted based on the GPU memory available
            enable_prefix_caching=self.enable_prefix_caching
        )

        self.tokenizer = self.llm.get_tokenizer()
        self.tokenizer.padding_side = "left"
        terminators = [
            self.tokenizer.eos_token_id,
            #self.tokenizer.convert_tokens_to_ids("
")
        ]

        self.json_schema = ClaimAssessment.model_json_schema()
        self.guided_decoding_params = GuidedDecodingParams(json=self.json_schema)

        self.sampling_params = SamplingParams(
            temperature=0.7,
            top_k=20,
            top_p=0.8,
            max_tokens=self.max_tokens,
            stop_token_ids=terminators,
            logprobs=5,
            guided_decoding=self.guided_decoding_params
        )


    def sent_tokenize_with_newlines(self, text):
        blocks = text.split('\n')
        
        tokenized_blocks = [sent_tokenize(block) for block in blocks]
        tokenized_text = []
        for block in tokenized_blocks:
            tokenized_text.extend(block)
            tokenized_text.append('\n')  

        return tokenized_text[:-1] 
    

    def apply_chat_template(self, doc, claim):

        user_prompt = self.user_prompt.replace("[DOCUMENT]", doc).replace("[CLAIM]", claim)
        message = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": user_prompt},
        ]
        text = self.tokenizer.apply_chat_template(message, add_generation_prompt=True, tokenize=False)

        return text

    def get_support_prob(self, response):
        """Extracts the probability of the first 'yes' token from all logprobs."""
        import math

        for token_logprobs in response.outputs[0].logprobs:
            for token_prob in token_logprobs.values():
                if token_prob.decoded_token.strip().lower() == 'yes':
                    return math.exp(token_prob.logprob) # Return on the first yes

        return 0
    
    def get_json_output(self, response):
        try:
            return json.loads(response.outputs[0].text)
        except json.JSONDecodeError:
            logger.warning(f"Failed to parse JSON from response")
        return {
            "claim_supported": "no",
            "need_calculation": "no",
            "calculation": ""
        }


    def get_all_chunks_per_doc(self, doc, claim):
        if doc in self.doc_chunk_cache:
            doc_chunks = self.doc_chunk_cache[doc]
        else:
            chonkie = ChonkieChunker(model_name=self.model_id, chunk_size=self.default_chunk_size)
            doc_chunks = [chunk.strip() for chunk in chonkie.chunk_text(doc) if chunk.strip()]
            self.doc_chunk_cache[doc] = doc_chunks

        if len(doc_chunks) == 0:
            doc_chunks = ['']

        claim_repeat = [claim] * len(doc_chunks)

        return {'doc_chunks': doc_chunks, 'claim_repeat': claim_repeat}


    def score(self, docs: List[str], claims: List[str], chunk_size=None) -> List[float]:

        self.doc_chunk_cache = {}
        self.chunk_size = chunk_size if chunk_size else self.default_chunk_size

        assert self.chunk_size < self.max_model_len, \
            "chunk_size must be less than max_model_len so that MiniCheck can process the claim"

        all_prompts = []
        doc_claim_indices = []
        for index, (doc, claim) in enumerate(zip(docs, claims)):
            chunks = self.get_all_chunks_per_doc(doc, claim)
            doc_chunks = chunks['doc_chunks']
            claim_repeat = chunks['claim_repeat']

            prompts = [self.apply_chat_template(doc_chunk, claim) for doc_chunk, claim in zip(doc_chunks, claim_repeat)]
            all_prompts.extend(prompts)
            doc_claim_indices.extend([index] * len(prompts))
        
        responses = self.llm.generate(
            all_prompts, 
            self.sampling_params,
        )

        json_outputs = [self.get_json_output(responses[idx]) for idx in range(len(responses))]

        probs_per_chunk = [self.get_support_prob(responses[idx]) for idx in range(len(responses))]

        result_dict = {}
        for index, (json_output, prob) in zip(doc_claim_indices, zip(json_outputs, probs_per_chunk)):
            if index not in result_dict:
                result_dict[index] = {
                    "claim_supported": False,
                    "need_calculation": False,
                    "calculations": [],
                    "support_probs": []
                }

            if json_output.get("claim_supported", "no").lower() == "yes":
                result_dict[index]["claim_supported"] = True

            if json_output.get("need_calculation", "no").lower() == "yes":
                result_dict[index]["need_calculation"] = True

            calc = json_output.get("calculation", "").strip()
            if calc and calc not in result_dict[index]["calculations"]:
                result_dict[index]["calculations"].append(calc)

            result_dict[index]["support_probs"].append(prob)


        probs_per_doc_claim_pair = [result_dict[index] for index in range(len(docs))] 
        pred_label, max_support_prob, used_chunk, support_prob_per_chunk = [], [], [], []
        need_calculation, calculations = [], []

        for idx in range(len(probs_per_doc_claim_pair)):
            doc = docs[idx]
            claim = claims[idx]
            result = probs_per_doc_claim_pair[idx]

            support_probs = result["support_probs"]
            is_supported = result["claim_supported"]
            needs_calc = result["need_calculation"]
            calc_list = result["calculations"]

            pred_label.append(1 if is_supported else 0)
            max_support_prob.append(max(support_probs) if support_probs else 0.0)
            used_chunk.append(self.get_all_chunks_per_doc(doc, claim)['doc_chunks'])
            support_prob_per_chunk.append(support_probs)
            need_calculation.append(needs_calc)
            calculations.append(calc_list)

        return pred_label, max_support_prob, need_calculation, calculations, used_chunk, support_prob_per_chunk