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

app = FastAPI()
API_KEY = os.environ.get("API_KEY")

try:
    model = AutoModelForCausalLM.from_pretrained("petertill/cordia-a6")
    tokenizer = AutoTokenizer.from_pretrained("petertill/cordia-a6")
    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
    print("Model and tokenizer loaded successfully!")

    class GenerateRequest(BaseModel):
        prompt: str
        key: str

    class GenerateResponse(BaseModel):
        generated_text: str

    @app.post("/generate", response_model=GenerateResponse)
    async def generate(request: GenerateRequest):
        authorization = request.key
        token = authorization.split('Bearer ')[1]
        if token != API_KEY:
            raise HTTPException(status_code=401, detail="Unauthorized")
        try:
            output = pipe(request.prompt, max_length=200)[0]['generated_text']
            return GenerateResponse(generated_text=output)
        except Exception as e:
            raise HTTPException(status_code=500, detail=str(e))

except Exception as e:
    print(f"Error: {e}")

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)