Update app.py

app.py

@@ -6,7 +6,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from openai import OpenAI
 
-print("Version ---- 4")
+print("Version ---- DeepSeek Only")
 app = FastAPI()
 
 # -----------------------------
@@ -15,7 +15,7 @@ app = FastAPI()
 class ConflictDetectionRequest(BaseModel):
     Req1: str
     Req2: str
-    model_choice: str   # "GPT-4", "DeepSeek-Reasoner", "LLaMA-3.1-8B-Instruct", "Fanar"
+    model_choice: str   # "GPT-4", "DeepSeek-Reasoner", "Fanar"
     prompt_type: str    # "zero-shot" or "few-shot"
     api_key: str = None # required only if model_choice == "GPT-4"
 
@@ -41,39 +41,21 @@ def build_prompt(req1, req2, prompt_type="zero-shot"):
         return f"Do the following sentences contradict each other, answer with just yes or no: 1.{req1} 2.{req2}"
 
 # -----------------------------
-# Startup: load models once
+# Startup: load DeepSeek once
 # -----------------------------
 @app.on_event("startup")
 def load_models():
-    print("Loading models into memory...")
-
-    # DeepSeek
+    print("Loading DeepSeek model into memory...")
     deepseek_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
     app.state.deepseek_tokenizer = AutoTokenizer.from_pretrained(deepseek_name)
     app.state.deepseek_tokenizer.pad_token = app.state.deepseek_tokenizer.eos_token
     app.state.deepseek_model = AutoModelForCausalLM.from_pretrained(
         deepseek_name,
-        dtype=torch.bfloat16,
-        device_map="auto"
+        torch_dtype=torch.float32   # CPU only
     )
 
-    # 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.")
-
 # -----------------------------
-# Model handlers (reuse loaded models)
+# Model handlers
 # -----------------------------
 def run_gpt4(req1, req2, prompt_type, api_key):
     client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=api_key)
@@ -95,7 +77,7 @@ def run_deepseek(req1, req2, prompt_type):
         return_tensors="pt",
         padding=True,
         truncation=True
-    ).to(model.device)
+    )
     outputs = model.generate(
         input_ids=inputs.input_ids,
         attention_mask=inputs.attention_mask,
@@ -104,24 +86,6 @@ 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_fanar(req1, req2, prompt_type):
     client = OpenAI(base_url="https://api.fanar.qa/v1", api_key=os.getenv("FANAR_API"))
     prompt = build_prompt(req1, req2, prompt_type)
@@ -145,11 +109,6 @@ 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 == "Fanar":
             answer = run_fanar(request.Req1, request.Req2, request.prompt_type)