Create app.py

app.py

@@ -0,0 +1,77 @@
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import os
+
+app = FastAPI(title="GPT-OSS-20B API")
+
+# Set environment variable for faster model downloads
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+
+# Model ID
+MODEL_ID = "openai/gpt-oss-20b"
+
+# Load tokenizer
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+
+# Load model with CPU offloading
+print("Loading model (this may take several minutes)...")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    device_map="auto",  # Automatically place on available devices (CPU)
+    torch_dtype="auto",  # Automatic precision
+    offload_folder="./offload",  # Offload weights to disk
+    max_memory={0: "15GB", "cpu": "30GB"},  # Memory constraints
+    trust_remote_code=True  # Required for custom model code
+)
+print(f"Model loaded on: {model.device}")
+print(f"Model dtype: {model.dtype}")
+
+# Enable gradient checkpointing to reduce memory usage
+model.gradient_checkpointing_enable()
+
+class ChatRequest(BaseModel):
+    message: str
+    max_tokens: int = 256
+    temperature: float = 0.7
+
+@app.post("/chat")
+async def chat_endpoint(request: ChatRequest):
+    try:
+        # Prepare input
+        messages = [{"role": "user", "content": request.message}]
+        inputs = tokenizer.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            return_dict=True
+        ).to("cpu")
+
+        # Generate response
+        with torch.no_grad():
+            generated = model.generate(
+                **inputs,
+                max_new_tokens=request.max_tokens,
+                temperature=request.temperature,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id,
+                repetition_penalty=1.1
+            )
+
+        # Decode response
+        response = tokenizer.decode(
+            generated[0][inputs["input_ids"].shape[-1]:],
+            skip_special_tokens=True
+        )
+        return {"response": response}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Clear cache regularly to manage memory
+torch.cuda.empty_cache()
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)