TobyYang7 commited on
Commit
a7191f1
1 Parent(s): 945f136

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +14 -11
app.py CHANGED
@@ -8,21 +8,24 @@ from PIL import Image
8
 
9
  import spaces
10
 
11
- # Model configuration
12
- model_id = "TheFinAI/FinLLaVA"
13
- device = "cuda:0"
14
- load_8bit = False
15
- load_4bit = False
 
 
 
 
16
 
17
- # Load the pretrained model
18
  tokenizer, llava_model, image_processor, context_len = load_pretrained_model(
19
- model_id,
20
  None,
21
  'llava_llama3',
22
- load_8bit,
23
- load_4bit,
24
- device=device
25
- )
26
 
27
  @spaces.GPU
28
  def bot_streaming(message, history):
 
8
 
9
  import spaces
10
 
11
+ parser = argparse.ArgumentParser()
12
+ parser.add_argument("--model-path", type=str, default="TheFinAI/FinLLaVA")
13
+ parser.add_argument("--device", type=str, default="cuda:0")
14
+ parser.add_argument("--conv-mode", type=str, default="llama_3")
15
+ parser.add_argument("--temperature", type=float, default=0)
16
+ parser.add_argument("--max-new-tokens", type=int, default=512)
17
+ parser.add_argument("--load-8bit", action="store_true")
18
+ parser.add_argument("--load-4bit", action="store_true")
19
+ args = parser.parse_args()
20
 
21
+ # load model
22
  tokenizer, llava_model, image_processor, context_len = load_pretrained_model(
23
+ args.model_path,
24
  None,
25
  'llava_llama3',
26
+ args.load_8bit,
27
+ args.load_4bit,
28
+ device=args.device)
 
29
 
30
  @spaces.GPU
31
  def bot_streaming(message, history):