xu3kev commited on
Commit
bb4bdd8
1 Parent(s): 6f70e6f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -17,7 +17,7 @@ MODEL_NAME="xu3kev/deepseekcoder-7b-logo-pbe"
17
  # MODEL_NAME="openlm-research/open_llama_3b"
18
  import torch
19
  from transformers import AutoModelForCausalLM, AutoTokenizer
20
- hug_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map='auto')
21
  hug_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
22
 
23
  INPUT_STRUCTION_TEMPLATE = """Here is a gray scale images representing with integer values 0-9.
@@ -223,7 +223,7 @@ def llm_call(question_prompt, model_name,
223
  top_p=1, n_samples=64, stop=None):
224
  if HUGGINGFACE:
225
  model_inputs = hug_tokenizer([question_prompt], return_tensors="pt").to('cuda')
226
- generated_ids = hug_model.generate(**model_inputs, max_length=1400, temperature=1, num_return_sequences=8, do_sample=True)
227
  responses = hug_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
228
  codes = []
229
  for response in responses:
 
17
  # MODEL_NAME="openlm-research/open_llama_3b"
18
  import torch
19
  from transformers import AutoModelForCausalLM, AutoTokenizer
20
+ hug_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map='auto', load_in_8bit=True)
21
  hug_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
22
 
23
  INPUT_STRUCTION_TEMPLATE = """Here is a gray scale images representing with integer values 0-9.
 
223
  top_p=1, n_samples=64, stop=None):
224
  if HUGGINGFACE:
225
  model_inputs = hug_tokenizer([question_prompt], return_tensors="pt").to('cuda')
226
+ generated_ids = hug_model.generate(**model_inputs, max_length=1400, temperature=1, num_return_sequences=32, do_sample=True)
227
  responses = hug_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
228
  codes = []
229
  for response in responses: