Update app.py

app.py

@@ -126,7 +126,7 @@ def evaluate(
     table,
     question,
     llm="alpaca-lora",
-    shot="1-shot",
+    num_shot="1-shot",
     input=None,
     temperature=0.1,
     top_p=0.75,
@@ -138,7 +138,7 @@ def evaluate(
     prompt_0shot = _INSTRUCTION + "\n" + _add_markup(table) + "\n" + "Q: " + question + "\n" + "A:"
     prompt = _TEMPLATE + "\n" + _add_markup(table) + "\n" + "Q: " + question + "\n" + "A:"
     if llm == "alpaca-lora":
-        if shot == "1-shot":
+        if num_shot == "1-shot":
             inputs = tokenizer(prompt, return_tensors="pt")
         else:
             inputs = tokenizer(prompt_0shot, return_tensors="pt")
@@ -161,7 +161,7 @@ def evaluate(
         s = generation_output.sequences[0]
         output = tokenizer.decode(s)
     elif llm == "flan-ul2":
-        if shot == "1-shot":
+        if num_shot == "1-shot":
             output = query({
                 "inputs": prompt
             })[0]["generated_text"]
@@ -179,18 +179,16 @@ def evaluate(
 model_deplot = Pix2StructForConditionalGeneration.from_pretrained("google/deplot", torch_dtype=torch.bfloat16).to(0)
 processor_deplot = Pix2StructProcessor.from_pretrained("google/deplot")
 
-def process_document(llm, image, question):
+def process_document(llm, num_shot, image, question):
     # image = Image.open(image)
     inputs = processor_deplot(images=image, text="Generate the underlying data table for the figure below:", return_tensors="pt").to(0, torch.bfloat16)
     predictions = model_deplot.generate(**inputs, max_new_tokens=512)
     table = processor_deplot.decode(predictions[0], skip_special_tokens=True).replace("<0x0A>", "\n")
 
     # send prompt+table to LLM
-    res = evaluate(table, question, llm=llm)
-    #return res + "\n\n" + res.split("A:")[-1]
+    res = evaluate(table, num_shot, question, llm=llm)
     if llm == "alpaca-lora":
         return [table, res.split("A:")[-1]]
-        # return [table, res]
     else:
         return [table, res]