Update app.py

app.py

@@ -1,58 +1,119 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
 from fastapi.responses import JSONResponse
+import os
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from accelerate import init_empty_weights, infer_auto_device_map, dispatch_model
+from openai import OpenAI
 
-print("Version ---- 1")
+print("Version ---- 2")
 app = FastAPI()
 
-# Load model and tokenizer from Hugging Face
-model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
+# -----------------------------
+# 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 # required only if model_choice == "GPT-4"
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# -----------------------------
+# 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}"
+    elif prompt_type == "few-shot":
+        # Example few-shot style (you can expand with more examples)
+        examples = (
+            "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}"
 
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    dtype=torch.bfloat16,
-    device_map="auto",
-    offload_folder="offload"   # folder for disk offload
-)
+# -----------------------------
+# Model handlers
+# -----------------------------
+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)
+    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):
+    model_name = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        dtype=torch.bfloat16,
+        device_map="auto"
+    )
+    prompt = build_prompt(req1, req2, prompt_type)
+    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
+    outputs = model.generate(inputs.input_ids, max_new_tokens=256)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-class ConflictDetectionRequest(BaseModel):
-    Req1: str
-    Req2: str
+def run_llama(req1, req2, prompt_type):
+    model_name = "meta-llama/Llama-3.1-8B-Instruct"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        dtype=torch.bfloat16,
+        device_map="auto"
+    )
+    prompt = build_prompt(req1, req2, prompt_type)
+    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
+    outputs = model.generate(inputs.input_ids, max_new_tokens=256)
+    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_KEY"))
+    prompt = build_prompt(req1, req2, prompt_type)
+    response = client.chat.completions.create(
+        model="Fanar",
+        messages=[{"role": "user", "content": prompt}]
+    )
+    return response.choices[0].message.content.strip()
 
+# -----------------------------
+# API route
+# -----------------------------
 @app.post("/predict")
 def predict(request: ConflictDetectionRequest):
-    Req1 = request.Req1  
-    Req2 = request.Req2
-
-    question = f"Do the following sentences contradict each other, answer with just yes or no: 1.{Req1}  2.{Req2}"
-    inputs = tokenizer([question], return_tensors="pt").to(model.device)
-    
-    # Generate response
-    outputs = model.generate(
-        input_ids=inputs.input_ids,
-        attention_mask=inputs.attention_mask,
-        max_new_tokens=512,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9
-    )
-    
-    # Decode and print response
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # print(response.split("### Response:")[-1].strip())
-    
-    content = {"resp": response.split("</think>")[1].strip(), "statusText": "OK","statusCode" : 0}
-
-    return JSONResponse(
-        content=content,
-        media_type="application/json",
-        status_code=200,
-    )
+    try:
+        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":
+            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)
+
+        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)