updated

app.py

@@ -13,7 +13,7 @@ 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"
 )
@@ -21,46 +21,23 @@ base_model = AutoModelForCausalLM.from_pretrained(
 # Load adapter and ensure it is on the correct device
 model = PeftModel.from_pretrained(base_model, adapter_path).to(device)
 
-# Load tokenizer and ensure padding token is set
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(base_model_path)
-tokenizer.pad_token = tokenizer.eos_token  # Avoids padding issues
 
 # Define request model for validation
 class GenerateRequest(BaseModel):
     prompt: str
 
-# **Use model.generate() instead of pipeline()**
+# **Use `model.generate()` instead of `pipeline()`**
 def generate_text_from_model(prompt: str):
     try:
-        input_data = tokenizer(
-            f"<s>[INST] {prompt} [/INST]", 
-            return_tensors="pt", 
-            padding=True, 
-            truncation=True
-        )
-        input_ids = input_data.input_ids.to(device)
-        attention_mask = input_data.attention_mask.to(device)
-
-        # Generate output
-        output_ids = model.generate(
-            input_ids, 
-            max_length=512, 
-            pad_token_id=tokenizer.eos_token_id,  # Explicitly setting pad_token_id
-            attention_mask=attention_mask
-        )
-
+        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)
-
-        # Extract only the assistant's response
-        response_text = generated_text.split("<|assistant|>\n")[-1].strip()
-        return response_text
-    except torch.cuda.OutOfMemoryError:
-        torch.cuda.empty_cache()
-        raise HTTPException(status_code=500, detail="CUDA Out of Memory. Try using a smaller model or lowering max_length.")
+        return generated_text
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-
 # Root endpoint for testing
 @app.get("/")
 async def root():