szymmon/SmolVLM_Essay_Knowledge_Distillation

1import gradio as gr2import torch3from transformers import AutoProcessor, Idefics3ForConditionalGeneration4import logging5 6logger = logging.getLogger(__name__)7 8class SimpleVLMInterface:9    def __init__(self):10        self.model = None11        self.processor = None12        self.initialize_model()13 14    def initialize_model(self):15        try:16            model_id = "HuggingFaceTB/SmolVLM-Instruct"17            self.model = Idefics3ForConditionalGeneration.from_pretrained(18                model_id,19                device_map="auto",20                torch_dtype=torch.bfloat1621            )22            self.processor = AutoProcessor.from_pretrained(model_id)23            # Load custom adapter24            adapter_path = "smolvlm-instruct-trl-sft-ChartQA"25            self.model.load_adapter(adapter_path)26        except Exception as e:27            logger.error(f"Error initializing model: {e}")28            raise29 30    def generate_response(31        self,32        text_input,33        image=None,34        max_tokens=512,35        temperature=0.7,36        top_p=0.9537    ):38        try:39            # Prepare the multimodal message format40            message_content = []41            42            # Add image content if provided43            if image is not None:44                if image.mode != 'RGB':45                    image = image.convert('RGB')46                message_content.append({47                    'type': 'image',48                    'image': image49                })50 51            # Add text content52            message_content.append({53                'type': 'text',54                'text': text_input55            })56 57            # Create the complete message structure58            messages = {59                'role': 'user',60                'content': message_content61            }62 63            # Apply chat template64            chat_input = self.processor.apply_chat_template(65                [messages],  # Wrap in list as it expects a sequence of messages66                add_generation_prompt=True67            )68 69            # Prepare model inputs70            model_inputs = self.processor(71                text=chat_input,72                images=[msg['image'] for msg in message_content if msg['type'] == 'image'] if image is not None else None,73                return_tensors="pt",74            ).to(self.model.device)75 76            # Generate response77            generated_ids = self.model.generate(78                **model_inputs,79                max_new_tokens=max_tokens,80                temperature=temperature,81                top_p=top_p,82                do_sample=True83            )84 85            # Process output86            trimmed_generated_ids = [87                out_ids[len(in_ids):] for in_ids, out_ids in zip(model_inputs.input_ids, generated_ids)88            ]89            output_text = self.processor.batch_decode(90                trimmed_generated_ids,91                skip_special_tokens=True,92                clean_up_tokenization_spaces=False93            )[0]94 95            return output_text96        except Exception as e:97            logger.error(f"Error generating response: {e}")98            return f"Error: {str(e)}"99 100def create_interface():101    vlm = SimpleVLMInterface()102    with gr.Blocks(title="Simple VLM Interface") as demo:103        with gr.Row():104            with gr.Column():105                image_input = gr.Image(type="pil", label="Upload Image (optional)")106                text_input = gr.Textbox(label="Enter your text", lines=2)107                with gr.Row():108                    max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max tokens")109                    temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")110                    top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p")111                submit_btn = gr.Button("Generate Response")112                output_text = gr.Textbox(label="Response", lines=4)113 114        submit_btn.click(115            fn=vlm.generate_response,116            inputs=[text_input, image_input, max_tokens, temperature, top_p],117            outputs=output_text118        )119 120    return demo121 122if __name__ == "__main__":123    demo = create_interface()124    demo.launch()