|
|
import os |
|
|
os.environ['TRANSFORMERS_CACHE'] = './cache' |
|
|
|
|
|
from fastapi import FastAPI |
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
|
|
|
app = FastAPI() |
|
|
|
|
|
model_name = "microsoft/BioGPT-Large" |
|
|
|
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name) |
|
|
model = AutoModelForCausalLM.from_pretrained(model_name) |
|
|
|
|
|
@app.get("/") |
|
|
def read_root(): |
|
|
return {"message": "BioGPT is ready!"} |
|
|
|
|
|
@app.post("/generate") |
|
|
def generate_text(prompt: str): |
|
|
inputs = tokenizer(prompt, return_tensors="pt") |
|
|
outputs = model.generate(**inputs, max_length=50) |
|
|
result = tokenizer.decode(outputs[0], skip_special_tokens=True) |
|
|
return {"response": result} |
|
|
|
