updated

app.py

@@ -2,51 +2,24 @@ from fastapi import FastAPI
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from peft import PeftModel
 import torch
-import os
 
 app = FastAPI()
 
-# Define paths
+# Load Model from Hugging Face Hub
 base_model_path = "NousResearch/Hermes-3-Llama-3.2-3B"
 adapter_path = "thinkingnew/llama_invs_adapter"
 
-# Check if GPU is available
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# Create offload folder if needed
-offload_dir = "./offload"
-os.makedirs(offload_dir, exist_ok=True)
-
-# Load base model with offloading support
 base_model = AutoModelForCausalLM.from_pretrained(
-    base_model_path,
-    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
-    device_map="auto",
-    offload_folder=offload_dir if device == "cpu" else None  # Offload to disk if running on CPU
+    base_model_path, torch_dtype=torch.float16, device_map="auto"
 )
-
-# Load adapter
 model = PeftModel.from_pretrained(base_model, adapter_path)
-
-# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(base_model_path)
 
-# Load pipeline once for better performance
-text_pipe = pipeline(
-    task="text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_length=512,
-    device=0 if device == "cuda" else -1  # Use GPU index 0 if available, otherwise CPU
-)
-
-# Root endpoint for testing
 @app.get("/")
 async def root():
     return {"message": "Model is running! Use /generate/ for text generation."}
-
-# Text generation endpoint
 @app.post("/generate/")
 async def generate_text(prompt: str):
-    result = text_pipe(f"<s>[INST] {prompt} [/INST]")
-    return {"response": result[0]['generated_text']}
+    pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer, max_length=512)
+    result = pipe(f"<s>[INST] {prompt} [/INST]")
+    return {"response": result[0]['generated_text']}