Update app.py

app.py

@@ -4,14 +4,22 @@ from arxiv2text import arxiv_to_text
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-def get_model(model_url="thepowerfuldeez/Qwen2-1.5B-Summarize"):
+def get_model(model_url="thepowerfuldeez/Qwen2-1.5B-Summarize", use_cpu=False):
     tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B-Instruct")
-    model = AutoModelForCausalLM.from_pretrained(
-        model_url,
-        bnb_4bit_compute_dtype=torch.bfloat16,
-        load_in_4bit=True,
-        attn_implementation="flash_attention_2",
-    )
+    if use_cpu:
+        model = OVModelForCausalLM.from_pretrained(
+            model_url,
+            device_map="cpu",
+            load_in_4bit=True,
+            attn_implementation="flash_attention_2",
+        )
+    else:
+        model = AutoModelForCausalLM.from_pretrained(
+            model_url,
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            load_in_4bit=True,
+            attn_implementation="flash_attention_2",
+        )
     return model, tokenizer
     
 
@@ -27,7 +35,7 @@ def call_llm(model, tokenizer, text):
     output = tokenizer.decode(new_tokens, skip_special_tokens=True)
     return output
 
-model, tokenizer = get_model()
+model, tokenizer = get_model(use_cpu=True)
 
 def summarize_pdf(pdf_url):
     extracted_text = arxiv_to_text(pdf_url)