Spaces:

sucheendra93
/

custom-model

No application file

Create app.py

098ac58 verified 10 months ago

547 Bytes

from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline

app = FastAPI()

generator = pipeline(
    "text-generation",
    model="meta-llama/Llama-3.1-8B-Instruct",  # or your HF model
    device=0  # remove if CPU
)

class Request(BaseModel):
    prompt: str
    max_new_tokens: int = 150

@app.post("/generate")
def generate(req: Request):
    output = generator(
        req.prompt,
        max_new_tokens=req.max_new_tokens,
        do_sample=True,
    )
    return {"result": output[0]["generated_text"]}

	from fastapi import FastAPI
	from pydantic import BaseModel
	from transformers import pipeline

	app = FastAPI()

	generator = pipeline(
	"text-generation",
	model="meta-llama/Llama-3.1-8B-Instruct", # or your HF model
	device=0 # remove if CPU
	)

	class Request(BaseModel):
	prompt: str
	max_new_tokens: int = 150

	@app.post("/generate")
	def generate(req: Request):
	output = generator(
	req.prompt,
	max_new_tokens=req.max_new_tokens,
	do_sample=True,
	)
	return {"result": output[0]["generated_text"]}