app.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import T5ForConditionalGeneration, T5Tokenizer
import torch

# Initialize FastAPI app
app = FastAPI()

# Load tokenizer and model from local directories
tokenizer = T5Tokenizer.from_pretrained("./tokenizer")  # Assumes tokenizer folder in the repo
model = T5ForConditionalGeneration.from_pretrained("./model")  # Assumes model folder in the repo
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()  # Set model to evaluation mode

# Define input structure
class QuestionGenerationRequest(BaseModel):
    context: str
    answer: str

# API endpoint for generating questions
@app.post("/generate_question")
async def generate_question(request: QuestionGenerationRequest):
    try:
        # Prepare input text
        text = f"context: {request.context} answer: {request.answer} </s>"

        # Tokenize the input text
        encoding = tokenizer.encode_plus(text, max_length=512, padding="max_length", return_tensors="pt")
        input_ids, attention_mask = encoding["input_ids"].to(device), encoding["attention_mask"].to(device)

        # Generate questions
        with torch.no_grad():
            beam_outputs = model.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                max_length=72,
                early_stopping=True,
                num_beams=5,
                num_return_sequences=3
            )

        # Decode the generated questions
        questions = [
            tokenizer.decode(beam_output, skip_special_tokens=True, clean_up_tokenization_spaces=True)
            for beam_output in beam_outputs
        ]

        # Return the generated questions
        return {"questions": questions}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# To run locally (if testing locally, not needed on Hugging Face Spaces)
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)