from fastapi import FastAPI, Request
from transformers import pipeline

app = FastAPI()

# Load the model
generator = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.1")

@app.get("/")
def read_root():
    return {"message": "Welcome to the free LLM API"}

@app.post("/generate")
async def generate_text(request: Request):
    data = await request.json()
    prompt = data.get("prompt", "")
    response = generator(prompt, max_length=200, do_sample=True)
    return {"response": response[0]['generated_text']}