app.py

from fastapi import FastAPI
import uvicorn

model_name = "DavidAU/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)