Osnly commited on
Commit
7509c8d
·
verified ·
1 Parent(s): ba5f231

Update src/visual_insight.py

Browse files
Files changed (1) hide show
  1. src/visual_insight.py +4 -2
src/visual_insight.py CHANGED
@@ -1,4 +1,6 @@
1
  from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
2
  import torch
3
  import json
4
  import re
@@ -10,8 +12,8 @@ model_id = "google/gemma-3n-E4B"
10
 
11
 
12
  cache_dir = "/tmp/hf_cache"
13
- tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token, cache_dir=cache_dir)
14
- model = AutoModelForCausalLM.from_pretrained(model_id, token=hf_token, cache_dir=cache_dir)
15
 
16
 
17
  def call_llm(prompt):
 
1
  from transformers import AutoModelForCausalLM, AutoTokenizer
2
+ from transformers import AutoProcessor, AutoModelForImageTextToText
3
+
4
  import torch
5
  import json
6
  import re
 
12
 
13
 
14
  cache_dir = "/tmp/hf_cache"
15
+ tokenizer = AutoProcessor.from_pretrained(model_id, token=hf_token, cache_dir=cache_dir)
16
+ model = AutoModelForImageTextToText.from_pretrained(model_id, token=hf_token, cache_dir=cache_dir)
17
 
18
 
19
  def call_llm(prompt):