Update app.py

app.py

@@ -7,6 +7,18 @@ import transformers
 from transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor
 from peft import PeftModel
 
+
+if torch.cuda.is_available():
+    device = "cuda"
+else:
+    device = "cpu"
+
+try:
+    if torch.backends.mps.is_available():
+        device = "mps"
+except:
+    pass
+
 ## CoT prompts
 
 def _add_markup(table):
@@ -24,7 +36,6 @@ def _add_markup(table):
         # just use the raw table if parsing fails
         return table
 
-
 _TABLE = """Year | Democrats | Republicans | Independents
 2004 | 68.1% | 45.0% | 53.0%
 2006 | 58.0% | 42.0% | 53.0%
@@ -36,7 +47,6 @@ _TABLE = """Year | Democrats | Republicans | Independents
 
 _INSTRUCTION = 'Read the table below to answer the following questions.'
 
-
 _TEMPLATE = f"""First read an example then the complete question for the second table.
 ------------
 {_INSTRUCTION}
@@ -56,7 +66,6 @@ A: Let's find the row of year 2007, that's Row 3. Let's extract the numbers on R
 
 ## alpaca-lora
 
-# debugging...
 assert (
     "LlamaTokenizer" in transformers._import_structure["models.llama"]
 ), "LLaMA is now in HuggingFace's main branch.\nPlease reinstall it: pip uninstall transformers && pip install git+https://github.com/huggingface/transformers.git"
@@ -67,17 +76,6 @@ tokenizer = LlamaTokenizer.from_pretrained("decapoda-research/llama-7b-hf")
 BASE_MODEL = "decapoda-research/llama-7b-hf"
 LORA_WEIGHTS = "tloen/alpaca-lora-7b"
 
-if torch.cuda.is_available():
-    device = "cuda"
-else:
-    device = "cpu"
-
-try:
-    if torch.backends.mps.is_available():
-        device = "mps"
-except:
-    pass
-
 if device == "cuda":
     model = LlamaForCausalLM.from_pretrained(
         BASE_MODEL,
@@ -147,7 +145,10 @@ def get_response_from_openai(prompt, model="gpt-3.5-turbo", max_output_tokens=25
   return ret
 
 ## deplot models
-model_deplot = Pix2StructForConditionalGeneration.from_pretrained("google/deplot", torch_dtype=torch.bfloat16).to(0)
+if device == "cuda":
+    model_deplot =  Pix2StructForConditionalGeneration.from_pretrained("google/deplot", torch_dtype=torch.bfloat16).to(0)
+else: 
+    model_deplot =  Pix2StructForConditionalGeneration.from_pretrained("google/deplot")
 processor_deplot = Pix2StructProcessor.from_pretrained("google/deplot")
 
 def evaluate(
@@ -199,7 +200,9 @@ def evaluate(
 
 def process_document(image, question, llm):
     # 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)
+    inputs = processor_deplot(images=image, text="Generate the underlying data table for the figure below:", return_tensors="pt")
+    if device == "cuda":
+        inputs = inputs.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")