Spaces:
Runtime error
Runtime error
File size: 895 Bytes
7431f37 74792b5 8d1870e 74792b5 c5712f4 74792b5 7431f37 8d1870e 7431f37 |
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 |
from fastapi import FastAPI, HTTPException
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import os
app = FastAPI()
# Hugging Face authentication token
hf_token = os.getenv("HF_TOKEN")
# Load model and tokenizer with the authentication token
model_name = "unsloth/codellama-34b-bnb-4bit"
tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=hf_token)
model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=hf_token)
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
@app.get("/")
def greet_json():
return {"Hello": "World!"}
@app.post("/generate")
def generate_text(prompt: str):
try:
generated_text = generator(prompt, max_length=50)
return {"generated_text": generated_text[0]['generated_text']}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
|