File size: 781 Bytes
0678348 be6e00e 0678348 1ed26cf 0678348 face0b6 1ed26cf 0678348 18044d1 0678348 9554a41 0678348 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 |
from fastapi import FastAPI
from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
app = FastAPI()
model_id = "ibleducation/ibl-fordham-7b-mistral"
tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=".")
model = AutoModelForCausalLM.from_pretrained(
model_id,
device_map="auto",
cache_dir=".",
low_cpu_mem_usage=True,
)
pipeline = pipeline(
"text-generation",
model=model,
tokenizer=tokenizer,
)
@app.get("/")
async def root():
return {"message": "Welcome to the Language Model API"}
@app.get("/generate/")
async def generate_text(prompt: str):
# Generate text based on the prompt received
prompt = f"<s>[INST]{prompt}</INST>"
response = pipeline(prompt)
return {"generated_text": response['generated_text']}
|