app.py

from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os

# Define the path where the model is located
model_directory = "tiny-gpt2"  # Make sure this is the correct relative or absolute path to your model folder

# Initialize the model and tokenizer using the correct directory path
tokenizer = AutoTokenizer.from_pretrained(model_directory)
model = AutoModelForCausalLM.from_pretrained(model_directory)

# FastAPI app
app = FastAPI()

class PromptRequest(BaseModel):
    prompt: str
    max_new_tokens: int = 50  # You can adjust the number of tokens generated

@app.post("/generate")
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}