Refactor model ID handling in app.py and update requirements.txt

app.py

@@ -22,44 +22,44 @@ def load_model_a(model_id):
     global tokenizer_a, model_a
     tokenizer_a = AutoTokenizer.from_pretrained(model_id)
     print(f"model A: {tokenizer_a.eos_token}")
-    model_a = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        attn_implementation="flash_attention_2",
-        trust_remote_code=True,
-    ).eval()
-    # try:
-    # except:
-    #     print(f"Using default attention implementation in {model_id}")
-    #     model_a = AutoModelForCausalLM.from_pretrained(
-    #         model_id,
-    #         torch_dtype=torch.bfloat16,
-    #         device_map="auto",
-    #         trust_remote_code=True,
-    #     ).eval()
+    try:
+        model_a = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            attn_implementation="flash_attention_2",
+            trust_remote_code=True,
+        ).eval()
+    except:
+        print(f"Using default attention implementation in {model_id}")
+        model_a = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True,
+        ).eval()
     return gr.update(label=model_id)
     
 def load_model_b(model_id):
     global tokenizer_b, model_b
     tokenizer_b = AutoTokenizer.from_pretrained(model_id)
     print(f"model B: {tokenizer_b.eos_token}")
-    model_b = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-        attn_implementation="flash_attention_2",
-        trust_remote_code=True,
-    ).eval()
-    # try:
-    # except:
-    #     print(f"Using default attention implementation in {model_id}")
-    #     model_b = AutoModelForCausalLM.from_pretrained(
-    #         model_id,
-    #         torch_dtype=torch.bfloat16,
-    #         device_map="auto",
-    #         trust_remote_code=True,
-    #     ).eval()
+    try:
+        model_b = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            attn_implementation="flash_attention_2",
+            trust_remote_code=True,
+        ).eval()
+    except:
+        print(f"Using default attention implementation in {model_id}")
+        model_b = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+            trust_remote_code=True,
+        ).eval()
     return gr.update(label=model_id)
 
 @spaces.GPU()