Update app.py

app.py

@@ -1,28 +1,37 @@
-import os
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-from langchain_community.llms import HuggingFacePipeline
-from langchain import PromptTemplate, LLMChain
+from transformers import (
+    pipeline,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    GenerationConfig,
+)
+from langchain_huggingface import HuggingFacePipeline
+from langchain_core.prompts import PromptTemplate
+from langchain.chains import LLMChain
 
-# — Model setup (small enough to CPU-serve in a Space) —
+# — 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)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True)
 
-# wrap in a HF pipeline and LangChain LLM
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
+# — Generation config & pipeline (new API) —
+gen_config = GenerationConfig(
     max_new_tokens=64,
     temperature=0.2,
     top_p=0.95,
     do_sample=False,
 )
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device_map="auto",
+    generation_config=gen_config,
+)
 llm = HuggingFacePipeline(pipeline=pipe)
 
-# define a simple prompt → chain
+# — Prompt & chain —
 prompt = PromptTemplate(
     input_variables=["description"],
     template=(
@@ -33,7 +42,7 @@ prompt = PromptTemplate(
 )
 chain = LLMChain(llm=llm, prompt=prompt)
 
-# FastAPI app
+# — FastAPI app —
 app = FastAPI()
 
 class Req(BaseModel):
@@ -44,7 +53,6 @@ class Res(BaseModel):
 
 @app.post("/generate-emmet", response_model=Res)
 async def generate_emmet(req: Req):
-    raw = chain.run(req.description)
-    # take just the first line after the prompt
+    raw = chain.invoke(req.description)  # use .invoke() instead of deprecated .run()
     emmet = raw.strip().splitlines()[0]
     return {"emmet": emmet}