Update app.py

app.py

@@ -1,29 +1,31 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-app = FastAPI()
-
-# Load model and tokenizer once at startup
-model_name = "./tiny-gpt2"  # ← path to the local directory
+# Initialize the model and tokenizer (Tiny GPT-2)
+model_name = "./tiny-gpt2"  # Path to your tiny-gpt2 folder
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
+# FastAPI app
+app = FastAPI()
 
 class PromptRequest(BaseModel):
     prompt: str
-    max_new_tokens: int = 50
+    max_new_tokens: int = 50  # You can adjust the number of tokens generated
 
 @app.post("/generate")
-async def generate_text(req: PromptRequest):
-    inputs = tokenizer(req.prompt, return_tensors="pt")
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=req.max_new_tokens,
-        do_sample=True,
-        temperature=0.8,
-        top_p=0.95,
-    )
-    generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"generated_text": generated}
+async def generate_text(request: PromptRequest):
+    # Encode the input prompt text
+    inputs = tokenizer.encode(request.prompt, return_tensors="pt")
+    
+    # Generate the text using the model
+    with torch.no_grad():
+        outputs = model.generate(inputs, max_length=request.max_new_tokens + len(inputs[0]))
+    
+    # Decode the generated text and return the response
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    return {"generated_text": generated_text}
+