File size: 1,027 Bytes
39cc8a5
823c760
a5b1f33
90b161e
97917f4
90b161e
1cadecd
 
 
 
0897499
1cadecd
90b161e
 
383a904
6d3fbf5
90b161e
 
 
 
 
 
 
1cadecd
90b161e
 
1cadecd
 
 
 
 
 
39cc8a5
 
 
 
 
 
a5b1f33
c6cb00e
65f1222
1cadecd
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
34
35
36
37
38
39
40
41
42
43
44
45
from fastapi import FastAPI
import uvicorn

model_name = "Llama-3.2-4X3B-MOE-Hell-California-Uncensored-10B-GGUF"

from transformers import AutoModel, AutoTokenizer, TextStreamer
import torch

# Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(
    model_name,
    device_map="auto",
    trust_remote_code=True
)

def llama2_chat(prompt):
    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.3
    )
    response = tokenizer.decode(output[0], skip_special_tokens=True)
    return response







app = FastAPI()

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

@app.get("/message")
async def message(input: str):
    return llama2_chat(input)

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