Update app.py

app.py

@@ -1,7 +1,30 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
 app = FastAPI()
 
+# Load the model and tokenizer
+model_name = "EleutherAI/gpt-neo-1.3B"  # Replace with your desired model
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
 @app.get("/")
 def read_root():
     return {"Hello": "World"}
+
+@app.post("/predict")
+async def predict(request: Request):
+    data = await request.json()
+    prompt = data.get("prompt", "")
+    if not prompt:
+        return {"error": "Prompt is required"}
+
+    # Tokenize the input
+    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")  # Use "cuda" if GPU is enabled
+
+    # Generate tokens
+    outputs = model.generate(inputs.input_ids, max_length=40, num_return_sequences=1)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    return {"response": response}