updated

app.py

@@ -15,7 +15,8 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load base model with device_map="auto" to handle GPUs automatically
 base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_path, torch_dtype=torch.float16, device_map="auto")
+    base_model_path, torch_dtype=torch.float16, device_map="auto"
+)
 
 # Load adapter and ensure it is on the correct device
 model = PeftModel.from_pretrained(base_model, adapter_path).to(device)
@@ -49,11 +50,7 @@ def generate_text_from_model(prompt: str):
         generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
         
         # Extract only the assistant's response
-        if "<|assistant|>" in generated_text:
-            response_text = generated_text.split("<|assistant|>")[-1].strip()
-        else:
-            response_text = generated_text.strip()
-        
+        response_text = generated_text.split("<|assistant|>\n")[-1].strip()
         return response_text
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -67,4 +64,4 @@ async def root():
 @app.post("/generate/")
 async def generate_text(request: GenerateRequest):
     response = generate_text_from_model(request.prompt)
-    return response
+    return {"response": response}