app.py

from fastapi import FastAPI
from pydantic import BaseModel
from transformers import (
    pipeline,
    AutoTokenizer,
    AutoModelForCausalLM,
)
from langchain_huggingface import HuggingFacePipeline
from langchain_core.prompts import PromptTemplate
from langchain.chains import LLMChain

# — Model setup —
MODEL_ID = "bigcode/starcoder2-3b"
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True)

# — Pipeline setup (pass generation parameters directly) —
pipe = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device_map="auto",
    max_new_tokens=64,
    temperature=0.2,
    top_p=0.95,
    do_sample=False,
)
llm = HuggingFacePipeline(pipeline=pipe)

# — Prompt & chain —
prompt = PromptTemplate(
    input_variables=["description"],
    template=(
        "### Convert English description to an Emmet abbreviation\n"
        "Description: {description}\n"
        "Emmet:"
    ),
)
chain = LLMChain(llm=llm, prompt=prompt)

# — FastAPI app —
app = FastAPI()

class Req(BaseModel):
    description: str

class Res(BaseModel):
    emmet: str

@app.post("/generate-emmet", response_model=Res)
async def generate_emmet(req: Req):
    raw = chain.invoke(req.description)  # use .invoke() instead of deprecated .run()
    emmet = raw.strip().splitlines()[0]
    return {"emmet": emmet}