Initial commit

app.py

@@ -0,0 +1,19 @@
+from fastapi import FastAPI, HTTPException
+from model import generate_code
+
+app = FastAPI()
+
+@app.get("/")
+def home():
+    return {"message": "Code Generation API is running!"}
+
+@app.post("/generate")
+def generate(prompt: str, max_tokens: int = 256):
+    if not prompt:
+        raise HTTPException(status_code=400, detail="Prompt cannot be empty.")
+
+    try:
+        code = generate_code(prompt, max_tokens)
+        return {"generated_code": code}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Internal Server Error: {str(e)}")

model.py

@@ -0,0 +1,16 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+# Load the pre-trained model
+MODEL_NAME = "bigcode/starcoder"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")
+
+def generate_code(prompt: str, max_tokens: int = 256):
+    """Generates code based on the input prompt."""
+    if not prompt.strip():
+        return "Error: Empty prompt provided."
+
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
+    output = model.generate(**inputs, max_new_tokens=max_tokens)
+    return tokenizer.decode(output[0], skip_special_tokens=True)

requirements.txt

@@ -0,0 +1,4 @@
+transformers
+torch
+fastapi
+uvicorn

runtime.txt

@@ -0,0 +1 @@
+python-3.9