Update app.py

app.py

@@ -6,7 +6,7 @@ import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from openai import OpenAI
 
-print("Version ---- 3")
+print("Version ---- 4")
 app = FastAPI()
 
 # -----------------------------
@@ -50,6 +50,7 @@ def load_models():
     # 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,
@@ -61,6 +62,7 @@ def load_models():
     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,
@@ -85,22 +87,39 @@ def run_gpt4(req1, req2, prompt_type, api_key):
     return completion.choices[0].message.content.strip()
 
 def run_deepseek(req1, req2, prompt_type):
-    print("Start run deepseek")
     tokenizer = app.state.deepseek_tokenizer
     model = app.state.deepseek_model
-    print("Start prompt building")
     prompt = build_prompt(req1, req2, prompt_type)
-    print("The prompt is ", prompt)
-    inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
-    outputs = model.generate(inputs.input_ids, max_new_tokens=256)
+    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").to(model.device)
-    outputs = model.generate(inputs.input_ids, max_new_tokens=256)
+    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):
@@ -125,7 +144,7 @@ def predict(request: ConflictDetectionRequest):
 
         elif request.model_choice == "DeepSeek-Reasoner":
             answer = run_deepseek(request.Req1, request.Req2, request.prompt_type)
-            print("Deepseek answer is : ", answer)
+
         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)