from fastapi import FastAPI
from pydantic import BaseModel

import torch
import os

from transformers import AutoTokenizer, AutoModelForCausalLM

# Fix missing dependency in HF Spaces
os.system("pip install tiktoken")

app = FastAPI()

MODEL_ID = "himalaya-ai/himalayagpt-0.5b"

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    use_fast=False
)

# Load model (safe settings for CPU HF Spaces)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    trust_remote_code=True,
    torch_dtype=torch.float32,
    low_cpu_mem_usage=True
)

class Request(BaseModel):
    prompt: str
    max_tokens: int = 100
    temperature: float = 0.2   # 🔥 FIXED (low = stable)


@app.get("/")
def home():
    return {"status": "running"}


@app.post("/generate")
def generate(req: Request):
    try:
        inputs = tokenizer(req.prompt, return_tensors="pt")

        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=req.max_tokens,

                do_sample=True,
                temperature=0.8,
                top_p=0.92,
                top_k=40,

                repetition_penalty=1.7,        # strong anti-loop
                no_repeat_ngram_size=5,        # stronger phrase blocking

                eos_token_id=tokenizer.eos_token_id,
                pad_token_id=tokenizer.eos_token_id
            )

        text = tokenizer.decode(
            outputs[0],
            skip_special_tokens=True
        )

        return {"response": text}

    except Exception as e:
        return {"error": str(e)}