Update app.py

app.py

@@ -10,8 +10,6 @@ from unsloth import FastLanguageModel
 import torch
 import re
 
-
-
 # Define helper functions
 async def fetch_data(url):
     headers = {
@@ -102,14 +100,13 @@ def translate_text(text):
         print(f"An error occurred during translation: {e}")
         return None
 
-@spaces.GPU()
-def summarize_url(url):
-
-    # Load the model
+def load_model():
     max_seq_length = 2048
     dtype = None
     load_in_4bit = True
-    
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
     model, tokenizer = FastLanguageModel.from_pretrained(
         model_name="unsloth/mistral-7b-instruct-v0.3-bnb-4bit",
         max_seq_length=max_seq_length,
@@ -117,9 +114,15 @@ def summarize_url(url):
         load_in_4bit=load_in_4bit,
     )
     
-    # Enable native 2x faster inference
-    FastLanguageModel.for_inference(model)
+    # Enable native 2x faster inference if GPU is available
+    if device == "cuda":
+        FastLanguageModel.for_inference(model)
     
+    return model, tokenizer, device
+
+model, tokenizer, device = load_model()
+
+def summarize_url(url):
     result = asyncio.run(fetch_data(url))
     text = concatenate_text(result)
     translated_text = translate_text(text)
@@ -136,7 +139,7 @@ def summarize_url(url):
     """
     
     prompt = alpaca_prompt.format(translated_text)
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
     outputs = model.generate(inputs.input_ids, max_new_tokens=64, use_cache=True)
     summary = tokenizer.decode(outputs[0], skip_special_tokens=True)