File size: 740 Bytes
7431f37
 
8d1870e
 
 
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
from fastapi import FastAPI, HTTPException
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

app = FastAPI()

# Load model and tokenizer
model_name = "meta-llama/Meta-Llama-3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
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))