import os

# ✅ Redirect HuggingFace cache to local folder
os.environ["TRANSFORMERS_CACHE"] = "./cache"
os.environ["HF_HOME"] = "./cache"  # just in case huggingface_hub also needs it

from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "Qwen/Qwen2.5-1.5B-Instruct"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)

device = 0 if torch.cuda.is_available() else -1
generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)

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

    prompt = f"""
You are an expert academic writer.

Task:
{user_input}

Instructions:
- Write clearly and professionally.
- Avoid repeating the question.
- Use academic tone.
- Format as a structured according to user requirements.
- Do not hallucinate facts.
"""

    result = generator(prompt, max_new_tokens=500, do_sample=False)[0]["generated_text"]

    if result.lower().startswith(prompt.lower()):
        result = result[len(prompt):].strip()

    return JSONResponse(content={"result": result})