hanzla commited on
Commit
c7721d9
1 Parent(s): 6c2bcf0

Update src/pdfchatbot.py

Browse files
Files changed (1) hide show
  1. src/pdfchatbot.py +2 -3
src/pdfchatbot.py CHANGED
@@ -105,9 +105,8 @@ class PDFChatBot:
105
  self.model = AutoModelForCausalLM.from_pretrained(
106
  self.config.get("autoModelForCausalLM"),
107
  device_map='auto',
108
- torch_dtype=torch.float32,
109
  token=True,
110
- load_in_8bit=False
111
  )
112
 
113
  def create_pipeline(self):
@@ -118,7 +117,7 @@ class PDFChatBot:
118
  model=self.model,
119
  task='text-generation',
120
  tokenizer=self.tokenizer,
121
- max_new_tokens=200
122
  )
123
  self.pipeline = HuggingFacePipeline(pipeline=pipe)
124
 
 
105
  self.model = AutoModelForCausalLM.from_pretrained(
106
  self.config.get("autoModelForCausalLM"),
107
  device_map='auto',
108
+ torch_dtype=torch.float16,
109
  token=True,
 
110
  )
111
 
112
  def create_pipeline(self):
 
117
  model=self.model,
118
  task='text-generation',
119
  tokenizer=self.tokenizer,
120
+ max_new_tokens=1024
121
  )
122
  self.pipeline = HuggingFacePipeline(pipeline=pipe)
123