Update app.py

app.py

@@ -55,22 +55,22 @@ def load_models():
         deepseek_name,
         dtype=torch.bfloat16,
         device_map="auto"
-    ).to("cuda" if torch.cuda.is_available() else "cpu")
+    )
 
     # LLaMA (requires HF_TOKEN secret)
-    llama_name = "meta-llama/Llama-3.1-8B-Instruct"
-    hf_token = os.getenv("LLAMA_HF_TOKEN")
-    if hf_token:
-        app.state.llama_tokenizer = AutoTokenizer.from_pretrained(llama_name, token=hf_token)
-        app.state.llama_tokenizer.pad_token = app.state.llama_tokenizer.eos_token
-        app.state.llama_model = AutoModelForCausalLM.from_pretrained(
-            llama_name,
-            token=hf_token,
-            dtype=torch.bfloat16,
-            device_map="auto"
-        )
-    else:
-        print("No HF_TOKEN found, LLaMA will not be available.")
+    # llama_name = "meta-llama/Llama-3.1-8B-Instruct"
+    # hf_token = os.getenv("LLAMA_HF_TOKEN")
+    # if hf_token:
+        # app.state.llama_tokenizer = AutoTokenizer.from_pretrained(llama_name, token=hf_token)
+        # app.state.llama_tokenizer.pad_token = app.state.llama_tokenizer.eos_token
+        # app.state.llama_model = AutoModelForCausalLM.from_pretrained(
+            # llama_name,
+            # token=hf_token,
+            # dtype=torch.bfloat16,
+            # device_map="auto"
+        # )
+    # else:
+        # print("No HF_TOKEN found, LLaMA will not be available.")
 
 # -----------------------------
 # Model handlers (reuse loaded models)
@@ -104,23 +104,23 @@ def run_deepseek(req1, req2, prompt_type):
     )
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-def run_llama(req1, req2, prompt_type):
-    tokenizer = app.state.llama_tokenizer
-    model = app.state.llama_model
-    prompt = build_prompt(req1, req2, prompt_type)
-    inputs = tokenizer(
-        [prompt],
-        return_tensors="pt",
-        padding=True,
-        truncation=True
-    ).to(model.device)
-    outputs = model.generate(
-        input_ids=inputs.input_ids,
-        attention_mask=inputs.attention_mask,
-        max_new_tokens=256,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+# def run_llama(req1, req2, prompt_type):
+    # tokenizer = app.state.llama_tokenizer
+    # model = app.state.llama_model
+    # prompt = build_prompt(req1, req2, prompt_type)
+    # inputs = tokenizer(
+        # [prompt],
+        # return_tensors="pt",
+        # padding=True,
+        # truncation=True
+    # ).to(model.device)
+    # outputs = model.generate(
+        # input_ids=inputs.input_ids,
+        # attention_mask=inputs.attention_mask,
+        # max_new_tokens=256,
+        # pad_token_id=tokenizer.eos_token_id
+    # )
+    # return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 def run_fanar(req1, req2, prompt_type):
     client = OpenAI(base_url="https://api.fanar.qa/v1", api_key=os.getenv("FANAR_API"))
@@ -145,10 +145,10 @@ def predict(request: ConflictDetectionRequest):
         elif request.model_choice == "DeepSeek-Reasoner":
             answer = run_deepseek(request.Req1, request.Req2, request.prompt_type)
 
-        elif request.model_choice == "LLaMA-3.1-8B-Instruct":
-            if not hasattr(app.state, "llama_model"):
-                return JSONResponse({"error": "LLaMA not loaded (missing HF_TOKEN)"}, status_code=400)
-            answer = run_llama(request.Req1, request.Req2, request.prompt_type)
+        # elif request.model_choice == "LLaMA-3.1-8B-Instruct":
+            # if not hasattr(app.state, "llama_model"):
+                # return JSONResponse({"error": "LLaMA not loaded (missing HF_TOKEN)"}, status_code=400)
+            # answer = run_llama(request.Req1, request.Req2, request.prompt_type)
 
         elif request.model_choice == "Fanar":
             answer = run_fanar(request.Req1, request.Req2, request.prompt_type)