Update app.py

app.py

@@ -44,14 +44,20 @@ if not torch.cuda.is_available():
 model = None
 tokenizer = None
 
-if torch.cuda.is_available():
-    model_id = "apple/OpenELM-3B-Instruct"    
-    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", trust_remote_code=True, low_cpu_mem_usage=True)
-    tokenizer_id = "meta-llama/Llama-2-7b-hf"
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
-    if tokenizer.pad_token == None:
-        tokenizer.pad_token = tokenizer.eos_token
-        tokenizer.pad_token_id = tokenizer.eos_token_id
+def initialize_model_and_tokenizer():
+    global model, tokenizer
+    if torch.cuda.is_available():
+        model_id = "apple/OpenELM-3B-Instruct"    
+        model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", trust_remote_code=True, low_cpu_mem_usage=True)
+        tokenizer_id = "meta-llama/Llama-2-7b-hf"
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+            tokenizer.pad_token_id = tokenizer.eos_token_id
+    else:
+        print("CUDA is not available. Model and tokenizer will not be initialized.")
+
+initialize_model_and_tokenizer()
 
 @spaces.GPU
 def generate(
@@ -63,36 +69,43 @@ def generate(
     top_k: int = 50,
     repetition_penalty: float = 1.4,
 ) -> Iterator[str]:
-    global model, tokenizer  # Access global variables
-
-    input_ids = tokenizer([message], return_tensors="pt").input_ids
-    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
-        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
-        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
-    input_ids = input_ids.to(model.device)
-
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-    generate_kwargs = dict(
-        {"input_ids": input_ids},
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        num_beams=1,
-        pad_token_id = tokenizer.eos_token_id,
-        repetition_penalty=repetition_penalty,
-        no_repeat_ngram_size=5,
-        early_stopping=True,
-    )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
+    global model, tokenizer
+
+    if tokenizer is None or model is None:
+        yield "Error: Model or tokenizer not initialized. Make sure you have GPU support and the necessary model access."
+        return
+
+    try:
+        input_ids = tokenizer([message], return_tensors="pt").input_ids
+        if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
+            input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
+            gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
+        input_ids = input_ids.to(model.device)
+
+        streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
+        generate_kwargs = dict(
+            input_ids=input_ids,
+            streamer=streamer,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            top_p=top_p,
+            top_k=top_k,
+            temperature=temperature,
+            num_beams=1,
+            pad_token_id=tokenizer.eos_token_id,
+            repetition_penalty=repetition_penalty,
+            no_repeat_ngram_size=5,
+            early_stopping=True,
+        )
+        t = Thread(target=model.generate, kwargs=generate_kwargs)
+        t.start()
+
+        outputs = []
+        for text in streamer:
+            outputs.append(text)
+            yield "".join(outputs)
+    except Exception as e:
+        yield f"An error occurred: {str(e)}"
 
 chat_interface = gr.ChatInterface(
     fn=generate,