Update app.py

app.py

@@ -1,109 +1,35 @@
-import os
-import torch
 from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
 
-from unsloth import FastLanguageModel
-from transformers import AutoTokenizer
-
-# -------------------------------
-# HF cache paths
-# -------------------------------
-os.environ["HF_HOME"] = "/tmp"
-os.environ["TRANSFORMERS_CACHE"] = "/tmp"
-
-# -------------------------------
-# FastAPI
-# -------------------------------
-
-os.environ["CUDA_VISIBLE_DEVICES"] = ""  # Force CPU
+app = FastAPI()
 
-import unsloth  # Now it won't try to use GPU
+# ---- Load your HF model repo ----
+MODEL_REPO = "hello-ram/mpt-model"
 
-app = FastAPI(title="Unsolth GPT OSS API")
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
+print("Loading model...")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_REPO,
+    torch_dtype=torch.float16,
+    device_map="auto"
 )
 
-# -------------------------------
-# Model variables
-# -------------------------------
-model = None
-tokenizer = None
-
-# Paths (exact as in your repo)
-base_model_name = "unsolth_gpt.20"  # your folder
-lora_model_path = "unsolth_gpt.20"  # LoRA files are inside same folder
-
-# -------------------------------
-# Load model
-# -------------------------------
-def load_model():
-    global model, tokenizer
-    if model is None or tokenizer is None:
-        tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
-
-        # Load base model on CPU
-        base_model = FastLanguageModel.from_pretrained(
-            base_model_name,
-            trust_remote_code=True,
-            device="cpu"
-        )
-
-        # Inject LoRA weights on CPU
-        model = FastLanguageModel.get_peft_model(
-            base_model,
-            r=8,
-            target_modules=[
-                "q_proj", "k_proj", "v_proj", "o_proj",
-                "gate_proj", "up_proj", "down_proj"
-            ],
-            lora_alpha=16,
-            lora_dropout=0,
-            bias="none",
-            state_dict=torch.load(os.path.join(lora_model_path, "model.safetensors"), map_location="cpu")
-        )
-        model.eval()
-
-# -------------------------------
-# Input schema
-# -------------------------------
-class QueryRequest(BaseModel):
-    question: str
-    max_new_tokens: int = 64
-    temperature: float = 0.7
-    top_p: float = 0.9
-    reasoning_effort: str = "medium"
-
-# -------------------------------
-# Health
-# -------------------------------
-@app.get("/")
-def health():
-    return {"status": "ok"}
-
-# -------------------------------
-# Predict
-# -------------------------------
-@app.post("/predict")
-def predict(req: QueryRequest):
-    load_model()
+class InputText(BaseModel):
+    text: str
 
-    inputs = tokenizer.apply_chat_template(
-        [{"role": "user", "content": req.question}],
-        add_generation_prompt=True,
-        return_tensors="pt",
-        return_dict=True,
-        reasoning_effort=req.reasoning_effort
-    ).to("cpu")  # force CPU
+@app.post("/generate")
+async def generate_text(data: InputText):
+    inputs = tokenizer(data.text, return_tensors="pt").to(model.device)
 
-    with torch.no_grad():
-        output = model.generate(**inputs, max_new_tokens=req.max_new_tokens)
+    output = model.generate(
+        **inputs,
+        max_new_tokens=200,
+        temperature=0.7
+    )
 
-    answer = tokenizer.decode(output[0], skip_special_tokens=True)
-    return {"question": req.question, "answer": answer}
+    generated = tokenizer.decode(output[0], skip_special_tokens=True)
+    return {"response": generated}