Update app.py

app.py

@@ -1,5 +1,4 @@
 from transformers import MllamaForConditionalGeneration, AutoProcessor, TextIteratorStreamer
-from transformers import StoppingCriteria, StoppingCriteriaList
 from PIL import Image
 import requests
 import torch
@@ -9,161 +8,12 @@ from gradio import FileData
 import time
 import spaces
 import re
-import copy
-
 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):
     
@@ -210,64 +60,20 @@ def bot_streaming(message, history, max_new_tokens=250):
 
     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)
+    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)  
+        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
+        yield buffer
 
 
 demo = gr.ChatInterface(fn=bot_streaming, title="LLaVA-CoT",
@@ -281,8 +87,7 @@ demo = gr.ChatInterface(fn=bot_streaming, title="LLaVA-CoT",
           )
         ],
       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]],
+                [{"text": "How to make this pastry?", "files": ["./Example2.png"]},512]],
       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",