AstraBert commited on
Commit
95b36ab
1 Parent(s): fa20ff3
Files changed (1) hide show
  1. utils.py +4 -5
utils.py CHANGED
@@ -1,4 +1,4 @@
1
- from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
2
  import time
3
  from langchain_community.llms import HuggingFacePipeline
4
  from langchain.storage import LocalFileStore
@@ -12,8 +12,8 @@ import os
12
  from pypdf import PdfMerger
13
 
14
 
15
- mod = "microsoft/phi-1_5"
16
- tsk = "text-generation"
17
 
18
  def merge_pdfs(pdfs: list):
19
  merger = PdfMerger()
@@ -136,10 +136,9 @@ try:
136
 
137
  tokenizer.pad_token = tokenizer.eos_token
138
 
139
- model = AutoModelForCausalLM.from_pretrained(
140
  mod,
141
  trust_remote_code=True,
142
- torch_dtype="auto",
143
  )
144
  except Exception as e:
145
  import sys
 
1
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
2
  import time
3
  from langchain_community.llms import HuggingFacePipeline
4
  from langchain.storage import LocalFileStore
 
12
  from pypdf import PdfMerger
13
 
14
 
15
+ mod = "google/flan-t5-large"
16
+ tsk = "text2text-generation"
17
 
18
  def merge_pdfs(pdfs: list):
19
  merger = PdfMerger()
 
136
 
137
  tokenizer.pad_token = tokenizer.eos_token
138
 
139
+ model = AutoModelForSeq2SeqLM.from_pretrained(
140
  mod,
141
  trust_remote_code=True,
 
142
  )
143
  except Exception as e:
144
  import sys