updated

app.py

@@ -1,4 +1,5 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
 import torch
@@ -12,23 +13,30 @@ adapter_path = "thinkingnew/llama_invs_adapter"
 # Check if GPU is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load base model with device_map="auto" to handle GPUs automatically
+# 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"
 )
 
-# Load adapter
-model = PeftModel.from_pretrained(base_model, adapter_path).to(device)  # Ensure model is on correct device
+# Load adapter and ensure it is on the correct device
+model = PeftModel.from_pretrained(base_model, adapter_path).to(device)
 
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(base_model_path)
 
+# Define request model for validation
+class GenerateRequest(BaseModel):
+    prompt: str
+
 # **Use `model.generate()` instead of `pipeline()`**
 def generate_text_from_model(prompt: str):
-    input_ids = tokenizer(f"<s>[INST] {prompt} [/INST]", return_tensors="pt").input_ids.to(device)
-    output_ids = model.generate(input_ids, max_length=512)
-    generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return generated_text
+    try:
+        input_ids = tokenizer(f"<s>[INST] {prompt} [/INST]", return_tensors="pt").input_ids.to(device)
+        output_ids = model.generate(input_ids, max_length=512)
+        generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        return generated_text
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 # Root endpoint for testing
 @app.get("/")
@@ -37,6 +45,6 @@ async def root():
 
 # Text generation endpoint
 @app.post("/generate/")
-async def generate_text(prompt: str):
-    response = generate_text_from_model(prompt)
+async def generate_text(request: GenerateRequest):
+    response = generate_text_from_model(request.prompt)
     return {"response": response}