Create app.py

app.py

@@ -0,0 +1,35 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+# Initialize FastAPI app
+app = FastAPI()
+
+# Load pre-trained DistilGPT-2 model and tokenizer
+model_name = "distilgpt2"  # Smaller GPT-2 model
+model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Pydantic model for request body
+class TextRequest(BaseModel):
+    text: str
+
+# Route to generate text
+@app.post("/generate/")
+async def generate_text(request: TextRequest):
+    # Encode the input text
+    inputs = tokenizer.encode(request.text, return_tensors="pt")
+
+    # Generate a response from the model
+    with torch.no_grad():
+        outputs = model.generate(inputs, max_length=100, num_return_sequences=1, no_repeat_ngram_size=2, top_p=0.9, top_k=50)
+
+    # Decode the generated response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return {"generated_text": response}
+
+# Optionally, you can add a root endpoint for checking server health
+@app.get("/")
+async def read_root():
+    return {"message": "Welcome to the GPT-2 FastAPI server!"}