File size: 919 Bytes
39cc8a5
823c760
97917f4
a5b1f33
97917f4
 
 
 
39cc8a5
 
 
 
 
 
a5b1f33
c6cb00e
65f1222
97917f4
 
 
 
 
 
 
 
 
 
 
8a5a310
 
823c760
73fcf85
 
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
32
33
from fastapi import FastAPI
import uvicorn
from transformers import AutoTokenizer, AutoModel

model_name = "TheBloke/Wizard-Vicuna-13B-Uncensored-GPTQ"

model = AutoModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

app = FastAPI()

@app.get("/")
def greet_json():
    return {"Hello": "World!"}

@app.get("/message")
async def message(input: str):
    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
    
    output = model.generate(
        input_ids=inputs["input_ids"],
        attention_mask=inputs["attention_mask"],  # Pass attention_mask!
        max_new_tokens=100,
        temperature=0.0,  # Disables randomness
        do_sample=False  # Greedy decoding
    )
    
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    
    return response

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)