app.py

from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer
from transformers import StoppingCriteria, StoppingCriteriaList
from PIL import Image
import requests
import torch
from threading import Thread
import gradio as gr
from gradio import FileData
import time
import spaces
import re
ckpt = "Xkev/Llama-3.2V-11B-cot"
model = MllamaForConditionalGeneration.from_pretrained(ckpt,
    torch_dtype=torch.bfloat16).to("cuda")
processor = AutoProcessor.from_pretrained(ckpt)


class StopOnStrings(StoppingCriteria):
    def __init__(self, stop_strings, tokenizer):
        self.stop_strings = stop_strings
        self.tokenizer = tokenizer

    def __call__(self, input_ids, scores, **kwargs):
        generated_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
        for stop_string in self.stop_strings:
            if stop_string in generated_text:
                return True
        return False
        
def judge(image, prompt, outputs, type="summary"):
    input_outputs = []
    kwargs = dict(do_sample=True, max_new_tokens=2048, temperature=0.6, top_p=0.9)
    
    hint = None
    if type == "all":
        judge_prompt = f'Now you act as a judge, helping me determine which of the two texts I provide better answers the question.'
        recall_prompt = ""
        for output in outputs:
            input_outputs.append(output)
    elif type == "sentence":
        judge_prompt = f'Now you act as a judge, helping me determine which of the two texts I provide is a better next sentence for the answer to the question.'
        recall_prompt = ""
        for output in outputs:
            sentences = output.split(".")
            if len(sentences) > 2:
                hint = ' '.join(sentences[:-2])
            input_outputs.append(sentences[-2])
    elif type == "summary":
        judge_prompt = f'Now you act as a judge, helping me determine which of the two texts I provide better provides a summary of what it should do to solve the question. The summary should focus on outlining the main approach instead of stating specific analytical reasoning or math formula.'
        recall_prompt = f'Please note that a better summary should focus on outlining the main approach instead of stating specific analytical reasoning or math formula.'
        for output in outputs:
            input_match = re.search(r'<SUMMARY>(.*?)</SUMMARY>', output, re.DOTALL)
            if input_match:
                input_outputs.append(input_match.group(1))
    elif type == "caption":
        judge_prompt = f'Now you act as a judge, helping me determine which of the two texts I provide better summarizes the information in the image related to the question, and has fewer errors. It is essential that the captions are as thorough as possible while remaining accurate, capturing as many details as possible rather than providing only general commentary.'
        recall_prompt = f'Please note that a better caption should be as thorough as possible while remaining accurate, capturing as many details as possible rather than providing only general commentary.'
        for output in outputs:
            input_match = re.search(r'<CAPTION>(.*?)</CAPTION>', output, re.DOTALL)
            if input_match:
                hint_match = re.search(r'<SUMMARY>(.*?)</SUMMARY>', output, re.DOTALL)
                if hint_match:
                    input_outputs.append(input_match.group(1))
    elif type == "reasoning":
        judge_prompt = f'Now you act as a judge, helping me determine which of the two texts I provide better explains the reasoning process to solve the question, and has fewer errors. Begin by thoroughly reviewing the question, followed by an in-depth examination of each answer individually, noting any differences. Subsequently, analyze these differences to determine which response demonstrates stronger reasoning and provide a clear conclusion.'
        recall_prompt = f'Begin by thoroughly reviewing the question, followed by an in-depth examination of each answer individually, noting any differences. Subsequently, analyze these differences to determine which response demonstrates stronger reasoning and provide a clear conclusion.'
        for output in outputs:
            input_match = re.search(r'<REASONING>(.*?)</REASONING>', output, re.DOTALL)
            if input_match:
                hint_match = re.search(r'<SUMMARY>(.*?)</SUMMARY>', output, re.DOTALL)
                if hint_match:
                    hint_caption_match = re.search(r'<CAPTION>(.*?)</CAPTION>', output, re.DOTALL)
                    if hint_caption_match:
                        hint = hint_caption_match.group(1)
                        input_outputs.append(input_match.group(1))
    elif type == "conclusion":
        judge_prompt = f'Now you act as a judge, helping me determine which of the two texts I provide offers a more effective conclusion to the question. The conclusion should align with the reasoning presented in the hint. The conclusion should never refuse to answer the question.'
        recall_prompt = f'Please note that a better conclusion should align with the reasoning presented in the hint. The conclusion should never refuse to answer the question.'
        for output in outputs:
            input_match = re.search(r'<CONCLUSION>(.*?)</CONCLUSION>', output, re.DOTALL)
            if input_match:
                hint_match = re.search(r'<SUMMARY>(.*?)</SUMMARY>', output, re.DOTALL)
                if hint_match:
                    hint_caption_match = re.search(r'<CAPTION>(.*?)</CAPTION>', output, re.DOTALL)
                    if hint_caption_match:
                        hint_reasoning_match = re.search(r'<REASONING>(.*?)</REASONING>', output, re.DOTALL)
                        if hint_reasoning_match:
                            hint = hint_caption_match.group(1) + hint_reasoning_match.group(1)
                            input_outputs.append(input_match.group(1))

    if type == "reasoning":
        reasoning_prompt = f"""Now you act as a judge, helping me determine whether the reasoning process in the given text is correct and accurate based on the given information.
        You should assume that the given information about the image is correct.
        You should only consider the reasoning process itself, not the correctness of the background information.  
        If the reasoning process invovles any calculations, you should verify the accuracy of the calculations.
        You should output 'correct' if you don't find any errors in the reasoning process, and 'incorrect' if you find any errors."""
        
        reasoning_prompt_1 = reasoning_prompt + f'\n\nGiven Information: {hint}' + f'\n\nReasoning Process: {input_outputs[0]}'
        reasoning_message_1 = [
            {'role': 'user', 'content': [
                {'type': 'text', 'text': reasoning_prompt_1}
            ]}
        ]
        reasoning_input_text_1 = processor.apply_chat_template(reasoning_message_1, add_generation_prompt=True)
        reasoning_inputs_1 = processor(None, reasoning_input_text_1, return_tensors='pt')
        reasoning_output_1 = model.generate(**reasoning_inputs_1, **kwargs)
        reasoning_output_text_1 = processor.decode(reasoning_output_1[0][reasoning_inputs_1['input_ids'].shape[1]:]).replace('<|eot_id|>', '').replace('<|endoftext|>', '')
        if "incorrect" in reasoning_output_text_1:
            #logging
            with open('log.jsonl', 'a') as f:
                json_obj = {
                    "prompt": prompt,
                    "outputs": outputs,
                    "judge_output": reasoning_output_text_1
                }
                f.write(json.dumps(json_obj) + '\n')
            return 1
        
        reasoning_prompt_2 = reasoning_prompt + f'\n\nGiven Information: {hint}' + f'\n\nReasoning Process: {input_outputs[1]}'
        reasoning_message_2 = [
            {'role': 'user', 'content': [
                {'type': 'text', 'text': reasoning_prompt_2}
            ]}
        ]
        reasoning_input_text_2 = processor.apply_chat_template(reasoning_message_2, add_generation_prompt=True)
        reasoning_inputs_2 = processor(None, reasoning_input_text_2, return_tensors='pt')
        reasoning_output_2 = model.generate(**reasoning_inputs_2, **kwargs)
        reasoning_output_text_2 = processor.decode(reasoning_output_2[0][reasoning_inputs_2['input_ids'].shape[1]:]).replace('<|eot_id|>', '').replace('<|endoftext|>', '')
        if "incorrect" in reasoning_output_text_2:
            #logging
            with open('log.jsonl', 'a') as f:
                json_obj = {
                    "prompt": prompt,
                    "outputs": outputs,
                    "judge_output": reasoning_output_text_2
                }
                f.write(json.dumps(json_obj) + '\n')
            return 0
            
    judge_prompt += f'\n\nQuestion: {prompt}'
    if hint:
        judge_prompt += f'\n\nHint about the Question: {hint}'
    for i, output in enumerate(input_outputs):
        judge_prompt += f'\nRepsonse {i+1}: {output}'
    judge_prompt += f'\n\n{recall_prompt}'
    judge_prompt += f' Please strictly follow the following format requirements when outputting, and don’t have any other unnecessary words.'
    judge_prompt += f'\n\nOutput format: "Since [reason], I choose response [1/2]."'
    
    judge_message = [
        {'role': 'user', 'content': [
            {'type': 'image'},
            {'type': 'text', 'text': judge_prompt}
        ]}
    ]
    judge_input_text = processor.apply_chat_template(judge_message, add_generation_prompt=True)
    judge_inputs = processor(image, judge_input_text, return_tensors='pt')
    judge_output = model.generate(**judge_inputs, **kwargs)
    judge_output_text = processor.decode(judge_output[0][judge_inputs['input_ids'].shape[1]:]).replace('<|eot_id|>', '').replace('<|endoftext|>', '')
    
    if "I choose response 1" in judge_output_text:
        return 0
    else:
        return 1
    
@spaces.GPU
def bot_streaming(message, history, max_new_tokens=250):
    
    txt = message["text"]
    ext_buffer = f"{txt}"
    
    messages= [] 
    images = []
    

    for i, msg in enumerate(history): 
        if isinstance(msg[0], tuple):
            messages.append({"role": "user", "content": [{"type": "text", "text": history[i+1][0]}, {"type": "image"}]})
            messages.append({"role": "assistant", "content": [{"type": "text", "text": history[i+1][1]}]})
            images.append(Image.open(msg[0][0]).convert("RGB"))
        elif isinstance(history[i-1], tuple) and isinstance(msg[0], str):
            # messages are already handled
            pass
        elif isinstance(history[i-1][0], str) and isinstance(msg[0], str): # text only turn
            messages.append({"role": "user", "content": [{"type": "text", "text": msg[0]}]})
            messages.append({"role": "assistant", "content": [{"type": "text", "text": msg[1]}]})

    # add current message
    if len(message["files"]) == 1:
        
        if isinstance(message["files"][0], str): # examples
            image = Image.open(message["files"][0]).convert("RGB")
        else: # regular input
            image = Image.open(message["files"][0]["path"]).convert("RGB")
        images.append(image)
        messages.append({"role": "user", "content": [{"type": "text", "text": txt}, {"type": "image"}]})
    else:
        messages.append({"role": "user", "content": [{"type": "text", "text": txt}]})


    texts = processor.apply_chat_template(messages, add_generation_prompt=True)

    if images == []:
        inputs = processor(text=texts, return_tensors="pt").to("cuda")
    else:
        inputs = processor(text=texts, images=images, return_tensors="pt").to("cuda")
        
    streamer = TextIteratorStreamer(processor, skip_special_tokens=True, skip_prompt=True)

    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=max_new_tokens, do_sample=True, temperature=0.6, top_p=0.9)
    generated_text = ""

    stages = ['<SUMMARY>', '<CAPTION>', '<REASONING>', '<CONCLUSION>']
    end_markers = ['</SUMMARY>', '</CAPTION>', '</REASONING>', '</CONCLUSION>']

    initial_length = len(inputs['input_ids'][0])
    input_ids = copy.deepcopy(inputs['input_ids'])

    for stage, end_marker in zip(stages, end_markers):
        stop_criteria = StoppingCriteriaList([StopOnStrings([end_marker], processor.tokenizer)])
        
        candidates = []
        for _ in range(2):  
            generation_kwargs.update({
                'stopping_criteria': stop_criteria
            })
            
            inputs = processor(image, input_ids, return_tensors='pt')
            output = model.generate(**inputs, **generation_kwargs)
            
            new_generated_ids = output[0]
            
            generated_text = processor.tokenizer.decode(new_generated_ids[initial_length:], skip_special_tokens=True)
            
            candidates.append({
                'input_ids': new_generated_ids.unsqueeze(0),
                'generated_text': generated_text,
            })
        
        while(len(candidates) > 1):
            candidate1 = candidates.pop(np.random.randint(len(candidates)))
            candidate2 = candidates.pop(np.random.randint(len(candidates)))
            outputs = [candidate1['generated_text'], candidate2['generated_text']]
            best_index = judge(image, prompt, outputs, type=stage[1:-1].lower())
            if best_index == 0:
                candidates.append(candidate1)
            else:
                candidates.append(candidate2)
        
        input_ids = candidates[0]['input_ids']

    final_output = processor.tokenizer.decode(input_ids[0][initial_length:], skip_special_tokens=True)
    final_output = re.sub(r"<(\w+)>", r"(Here begins the \1 stage)", final_output) 
    final_output = re.sub(r"</(\w+)>", r"(Here ends the \1 stage)", final_output) 
    return final_output

    # thread = Thread(target=model.generate, kwargs=generation_kwargs)
    # thread.start()
    # buffer = ""
    
    # for new_text in streamer:
    #     buffer += new_text
    #     generated_text_without_prompt = buffer
    #     time.sleep(0.01)
    
    #     buffer = re.sub(r"<(\w+)>", r"(Here begins the \1 stage)", buffer)  
    #     buffer = re.sub(r"</(\w+)>", r"(Here ends the \1 stage)", buffer)  
    
    #     yield buffer


demo = gr.ChatInterface(fn=bot_streaming, title="LLaVA-CoT",
      textbox=gr.MultimodalTextbox(), 
      additional_inputs = [gr.Slider(
              minimum=512,
              maximum=1024,
              value=512,
              step=1,
              label="Maximum number of new tokens to generate",
          )
        ],
      examples=[[{"text": "What is on the flower?", "files": ["./Example1.webp"]},512],
                [{"text": "How to make this pastry?", "files": ["./Example2.png"]},512],
                [{"text": f"Hint: Please answer the question and provide the correct option letter, e.g., A, B, C, D, at the end.\n Question: Subtract all tiny shiny balls. Subtract all purple objects. How many objects are left?\n Options:\n A. 4\n B. 8\n C. 2\n D. 6", "files": ["./reasoning.png"]},2048]],
      cache_examples=False,
      description="Upload an image, and start chatting about it. To learn more about LLaVA-CoT, visit [our GitHub page](https://github.com/PKU-YuanGroup/LLaVA-CoT). Note: Since Gradio currently does not support displaying the special markings in the output, we have replaced it with the expression (Here begins the X phase).",
      stop_btn="Stop Generation", 
      fill_height=True,
    multimodal=True)
    
demo.launch(debug=True)