Update app.py

app.py

@@ -6,74 +6,121 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from openai import OpenAI
 
+print("Version ---- 4")
 app = FastAPI()
 
-# Globals for models
-deepseek_model = None
-deepseek_tokenizer = None
-llama_model = None
-llama_tokenizer = None
-
+# -----------------------------
+# Request schema
+# -----------------------------
 class ConflictDetectionRequest(BaseModel):
     Req1: str
     Req2: str
     model_choice: str   # "GPT-4", "DeepSeek-Reasoner", "LLaMA-3.1-8B-Instruct", "Fanar"
     prompt_type: str    # "zero-shot" or "few-shot"
-    api_key: str = None
+    api_key: str = None # required only if model_choice == "GPT-4"
 
+# -----------------------------
+# Prompt builder
+# -----------------------------
 def build_prompt(req1, req2, prompt_type="zero-shot"):
     if prompt_type == "zero-shot":
         return f"Do the following sentences contradict each other, answer with just yes or no: 1.{req1} 2.{req2}"
-    else:
+    elif prompt_type == "few-shot":
         examples = (
-            "Example 1:\nReq1: The system shall allow password reset.\nReq2: The system shall not allow password reset.\nAnswer: yes\n\n"
-            "Example 2:\nReq1: The system shall support Arabic language.\nReq2: The system shall support English language.\nAnswer: no\n\n"
+            "Example 1:\n"
+            "Req1: The system shall allow password reset.\n"
+            "Req2: The system shall not allow password reset.\n"
+            "Answer: yes\n\n"
+            "Example 2:\n"
+            "Req1: The system shall support Arabic language.\n"
+            "Req2: The system shall support English language.\n"
+            "Answer: no\n\n"
         )
         return examples + f"Now answer: Do the following sentences contradict each other? 1.{req1} 2.{req2}"
+    else:
+        return f"Do the following sentences contradict each other, answer with just yes or no: 1.{req1} 2.{req2}"
 
-def run_deepseek(req1, req2, prompt_type):
-    global deepseek_model, deepseek_tokenizer
-    if deepseek_model is None:
-        print("Loading DeepSeek model into memory...")
-        model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
-        deepseek_tokenizer = AutoTokenizer.from_pretrained(model_name)
-        deepseek_tokenizer.pad_token = deepseek_tokenizer.eos_token
-        deepseek_model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            torch_dtype=torch.float32   # CPU only
+# -----------------------------
+# Startup: load models once
+# -----------------------------
+@app.on_event("startup")
+def load_models():
+    print("Loading models into memory...")
+
+    # DeepSeek
+    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"
+    ).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.")
+
+# -----------------------------
+# Model handlers (reuse loaded models)
+# -----------------------------
+def run_gpt4(req1, req2, prompt_type, api_key):
+    client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=api_key)
     prompt = build_prompt(req1, req2, prompt_type)
-    inputs = deepseek_tokenizer([prompt], return_tensors="pt", padding=True, truncation=True)
-    outputs = deepseek_model.generate(
+    completion = client.chat.completions.create(
+        model="openai/gpt-4",
+        messages=[{"role": "user", "content": prompt}],
+        temperature=0.7,
+        max_tokens=512
+    )
+    return completion.choices[0].message.content.strip()
+
+def run_deepseek(req1, req2, prompt_type):
+    tokenizer = app.state.deepseek_tokenizer
+    model = app.state.deepseek_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=deepseek_tokenizer.eos_token_id
+        pad_token_id=tokenizer.eos_token_id
     )
-    return deepseek_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 def run_llama(req1, req2, prompt_type):
-    global llama_model, llama_tokenizer
-    if llama_model is None:
-        print("Loading LLaMA model into memory...")
-        model_name = "meta-llama/Llama-3.1-8B-Instruct"
-        hf_token = os.getenv("LLAMA_HF_TOKEN")
-        llama_tokenizer = AutoTokenizer.from_pretrained(model_name, token=hf_token)
-        llama_tokenizer.pad_token = llama_tokenizer.eos_token
-        llama_model = AutoModelForCausalLM.from_pretrained(
-            model_name,
-            token=hf_token,
-            torch_dtype=torch.float32   # CPU only
-        )
+    tokenizer = app.state.llama_tokenizer
+    model = app.state.llama_model
     prompt = build_prompt(req1, req2, prompt_type)
-    inputs = llama_tokenizer([prompt], return_tensors="pt", padding=True, truncation=True)
-    outputs = llama_model.generate(
+    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=llama_tokenizer.eos_token_id
+        pad_token_id=tokenizer.eos_token_id
     )
-    return llama_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    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"))
@@ -84,30 +131,32 @@ def run_fanar(req1, req2, prompt_type):
     )
     return response.choices[0].message.content.strip()
 
+# -----------------------------
+# API route
+# -----------------------------
 @app.post("/predict")
 def predict(request: ConflictDetectionRequest):
     try:
-        if request.model_choice == "DeepSeek-Reasoner":
+        if request.model_choice == "GPT-4":
+            if not request.api_key:
+                return JSONResponse({"error": "API key required for GPT-4"}, status_code=400)
+            answer = run_gpt4(request.Req1, request.Req2, request.prompt_type, request.api_key)
+
+        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)
-        elif request.model_choice == "GPT-4":
-            if not request.api_key:
-                return JSONResponse({"error": "API key required for GPT-4"}, status_code=400)
-            client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=request.api_key)
-            prompt = build_prompt(request.Req1, request.Req2, request.prompt_type)
-            completion = client.chat.completions.create(
-                model="openai/gpt-4",
-                messages=[{"role": "user", "content": prompt}],
-                temperature=0.7,
-                max_tokens=512
-            )
-            answer = completion.choices[0].message.content.strip()
+
         else:
             return JSONResponse({"error": "Invalid model_choice"}, status_code=400)
 
         return JSONResponse({"resp": answer, "statusText": "OK", "statusCode": 0}, status_code=200)
+
     except Exception as e:
         return JSONResponse({"error": str(e)}, status_code=500)