Added app.py and requirements.txt

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from peft import PeftModel
@@ -7,44 +7,44 @@ import os
 
 app = FastAPI()
 
-# Load Hugging Face token from environment
-HF_TOKEN = os.environ.get("HF_TOKEN")
-
-# Base and LoRA model names
+HF_TOKEN = os.getenv("HF_TOKEN")
 BASE_MODEL = "google/gemma-2b-it"
 LORA_MODEL = "varshithkumar/gemma-finetuned-sql"
 
-# Load base model and tokenizer
+# Choose device
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 print("Loading base model...")
 base_model = AutoModelForCausalLM.from_pretrained(
     BASE_MODEL,
-    device_map="auto",
-    token=HF_TOKEN
-)
+    device_map=None,  # Avoid auto offloading
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    use_auth_token=HF_TOKEN
+).to(device)
+
+print("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(
     BASE_MODEL,
     use_fast=True,
-    token=HF_TOKEN
+    use_auth_token=HF_TOKEN
 )
 
-# Apply LoRA weights
 print("Applying LoRA adapter...")
 model = PeftModel.from_pretrained(
     base_model,
     LORA_MODEL,
-    token=HF_TOKEN
+    use_auth_token=HF_TOKEN
 )
 
 print("Model loaded successfully!")
 
-# Define input schema
 class InputData(BaseModel):
     prompt: str
     max_length: int = 100
 
 @app.post("/generate")
 def generate_text(data: InputData):
-    inputs = tokenizer(data.prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(data.prompt, return_tensors="pt").to(device)
     outputs = model.generate(**inputs, max_length=data.max_length)
     text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {"response": text}

requirements.txt

@@ -1,9 +1,7 @@
 fastapi
 uvicorn
 transformers
-peft
-torch
 accelerate
+peft
 bitsandbytes
 sentencepiece
-huggingface_hub