from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer

# Define the input schema
class ModelInput(BaseModel):
    prompt: str
    max_new_tokens: int = 50  # Optional: Defaults to 50 tokens

# Initialize FastAPI app
app = FastAPI()

# Load your model and tokenizer
model_path = "khurrameycon/SmolLM-135M-Instruct-qa_pairs_converted.json-25epochs"  # Update with your model directory
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path)

# Initialize the pipeline
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

@app.post("/generate")
def generate_text(input: ModelInput):
    try:
        result = generator(
            input.prompt,
            max_new_tokens=input.max_new_tokens,
            return_full_text=False,
        )
        return {"generated_text": result[0]["generated_text"]}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
def root():
    return {"message": "Welcome to the Hugging Face Model API!"}