saylee-m commited on
Commit
e0a50f9
1 Parent(s): b0ab786

added comments

Browse files
Files changed (1) hide show
  1. app.py +9 -6
app.py CHANGED
@@ -44,13 +44,13 @@ def load_models():
44
  # load donut
45
  donut_model, donut_processor = load_donut_model()
46
  print("donut downloaded")
47
- #load paligemma
48
- pg_model, pg_processor = load_paligemma_docvqa()
49
- print("paligemma downloaded")
50
 
51
  return {"donut":[donut_model, donut_processor],
52
- # "idefics": [idf_model, idf_processor],
53
- "paligemma": [pg_model, pg_processor]}
54
 
55
  loaded_models = load_models()
56
 
@@ -112,7 +112,7 @@ def process_document_pg(image_array, question):
112
 
113
  inputs = processor(images=image_array, text=question, return_tensors="pt").to(device)
114
  predictions = model.generate(**inputs, max_new_tokens=100)
115
- return processor.batch_decode(predictions[0], skip_special_tokens=True)[len(question):].lstrip("\n")
116
 
117
  def process_document_idf(image_array, question):
118
  model, processor = loaded_models.get("idefics")
@@ -124,6 +124,7 @@ def process_document_idf(image_array, question):
124
 
125
  def generate_answer_donut(image_array, question):
126
  try:
 
127
  answer = process_document_donut(image_array, question)
128
  print(answer)
129
  return answer
@@ -134,6 +135,7 @@ def generate_answer_donut(image_array, question):
134
 
135
  def generate_answer_idefics(image_array, question):
136
  try:
 
137
  # answer = process_document_idf(image_array, question)
138
  answer = inference_calling_idefics(image_array, question)
139
  print(answer)
@@ -145,6 +147,7 @@ def generate_answer_idefics(image_array, question):
145
 
146
  def generate_answer_paligemma(image_array, question):
147
  try:
 
148
  answer = process_document_pg(image_array, question)
149
  print(answer)
150
  return answer
 
44
  # load donut
45
  donut_model, donut_processor = load_donut_model()
46
  print("donut downloaded")
47
+ # #load paligemma
48
+ # pg_model, pg_processor = load_paligemma_docvqa()
49
+ # print("paligemma downloaded")
50
 
51
  return {"donut":[donut_model, donut_processor],
52
+ # "paligemma": [pg_model, pg_processor]
53
+ }
54
 
55
  loaded_models = load_models()
56
 
 
112
 
113
  inputs = processor(images=image_array, text=question, return_tensors="pt").to(device)
114
  predictions = model.generate(**inputs, max_new_tokens=100)
115
+ return processor.decode(predictions[0], skip_special_tokens=True)[len(question):].lstrip("\n")
116
 
117
  def process_document_idf(image_array, question):
118
  model, processor = loaded_models.get("idefics")
 
124
 
125
  def generate_answer_donut(image_array, question):
126
  try:
127
+ print("processing document - donut")
128
  answer = process_document_donut(image_array, question)
129
  print(answer)
130
  return answer
 
135
 
136
  def generate_answer_idefics(image_array, question):
137
  try:
138
+ print("processing document - idf2")
139
  # answer = process_document_idf(image_array, question)
140
  answer = inference_calling_idefics(image_array, question)
141
  print(answer)
 
147
 
148
  def generate_answer_paligemma(image_array, question):
149
  try:
150
+ print("processing document - pg")
151
  answer = process_document_pg(image_array, question)
152
  print(answer)
153
  return answer