Update src/insight.py

src/insight.py

@@ -1,4 +1,6 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoProcessor, AutoModelForImageTextToText
+
 import torch
 import json
 import os
@@ -11,8 +13,8 @@ model_id = "google/gemma-3n-E4B"
 
 hf_token = os.environ.get("HUGGINGFACE_TOKEN")
 cache_dir = "/tmp/hf_cache"
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token, cache_dir=cache_dir)
-model = AutoModelForCausalLM.from_pretrained(model_id, token=hf_token, cache_dir=cache_dir)
+tokenizer = AutoProcessor.from_pretrained(model_id, token=hf_token, cache_dir=cache_dir)
+model = AutoModelForImageTextToText.from_pretrained(model_id, token=hf_token, cache_dir=cache_dir)
 
 
 def call_llm(prompt):