import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)

class RequestData(BaseModel):
    inputs: str


def generate_text(prompt):

    # ✅ Proper chat formatting (THIS IS THE FIX)
    formatted_prompt = f"<|user|>\n{prompt}\n<|assistant|>\n"

    inputs = tokenizer(formatted_prompt, return_tensors="pt")

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=250,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id
        )

    result = tokenizer.decode(output[0], skip_special_tokens=True)

    # ✅ Extract only assistant response
    if "<|assistant|>" in result:
        result = result.split("<|assistant|>")[-1]

    return result.strip()


@app.post("/generate")
async def generate(request: RequestData):

    text = generate_text(request.inputs)

    return {
        "data": [text]
    }