from fastapi import FastAPI, Request
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import torch
import uvicorn

app = FastAPI(title="Phi-2 Free API")

# ✅ Non-gated model (public, CPU-friendly)
model_name = "microsoft/phi-2"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)

pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    max_new_tokens=200,
    temperature=0.7,
    top_p=0.9,
    repetition_penalty=1.1,
)

@app.get("/")
async def home():
    return {"message": "✅ Phi-2 API is running!"}

@app.post("/chat")
async def chat(request: Request):
    data = await request.json()
    prompt = data.get("prompt", "").strip()

    if not prompt:
        return {"error": "Please provide a 'prompt' in JSON."}

    # Generate text
    result = pipe(prompt, max_new_tokens=150, do_sample=True)[0]["generated_text"]

    # ✂️ Clean response: remove prompt echo
    cleaned = result.replace(prompt, "").strip()

    if not cleaned:
        cleaned = "(No response generated. Try a longer prompt or lower temperature.)"

    return {"response": cleaned}

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