Spaces:
Sleeping
Sleeping
File size: 2,117 Bytes
1fceb40 c7e6bba 1fceb40 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | from fastapi import FastAPI
from pydantic import BaseModel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
app = FastAPI(
title="Rapnss DevOps-Ultra-125M API",
description="A pure REST API for generating Python code via HTTPS using the custom Rapnss architecture",
version="1.0.0"
)
# Configuration
REPO_ID = "Rapnss/DevOps-Ultra-125M"
DEVICE = "cpu"
model = None
tokenizer = None
@app.on_event("startup")
def load_model():
global model, tokenizer
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(REPO_ID, trust_remote_code=True)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
print("Loading custom Rapnss model weights directly from the Hub...")
# trust_remote_code=True automatically downloads your custom PyTorch architecture files!
model = AutoModelForCausalLM.from_pretrained(REPO_ID, trust_remote_code=True)
model.to(DEVICE)
model.eval()
print("API Ready for Requests!")
class CodeRequest(BaseModel):
instruction: str
input_code: str = ""
max_tokens: int = 150
temperature: float = 0.7
@app.post("/v1/generate")
def generate_code(req: CodeRequest):
if model is None or tokenizer is None:
return {"error": "Model not loaded yet."}
prompt = f"Instruction: {req.instruction}\n"
if req.input_code:
prompt += f"Input: {req.input_code}\n"
prompt += "Output:\n"
input_ids = tokenizer.encode(prompt, return_tensors="pt").to(DEVICE)
output_ids = model.generate(
input_ids,
max_new_tokens=req.max_tokens,
do_sample=False,
pad_token_id=tokenizer.eos_token_id
)
response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
# Clean up the output to only return the generated code
if "Output:\n" in response:
response = response.split("Output:\n")[-1]
return {"generated_code": response.strip()}
@app.get("/")
def health_check():
return {"status": "active", "model": "Rapnss DevOps-Ultra-125M API is running!"}
|