oliver-aizip commited on
Commit
e6127a4
·
1 Parent(s): b8ee0a2
Files changed (1) hide show
  1. utils/models.py +2 -2
utils/models.py CHANGED
@@ -170,8 +170,8 @@ def run_inference(model_name, context, question, result_queue):
170
  # max_length=2048, # Keep original max_length for now
171
  # add_generation_prompt=True,
172
  # ).to(device)
173
- output = pipe(text_input, max_new_tokens=512)
174
- result = output[0]['generated_text'][-1]['content']
175
  # # Ensure input does not exceed model max length after adding generation prompt
176
  # # This check might be redundant if tokenizer handles it, but good for safety
177
  # # if actual_input.shape[1] > tokenizer.model_max_length:
 
170
  # max_length=2048, # Keep original max_length for now
171
  # add_generation_prompt=True,
172
  # ).to(device)
173
+ outputs = pipe(text_input, max_new_tokens=512)
174
+ result = outputs[0]['generated_text'][-1]['content']
175
  # # Ensure input does not exceed model max length after adding generation prompt
176
  # # This check might be redundant if tokenizer handles it, but good for safety
177
  # # if actual_input.shape[1] > tokenizer.model_max_length: