Spaces:
Sleeping
Sleeping
| # app.py | |
| import torch | |
| import uvicorn | |
| from fastapi import FastAPI | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from pydantic import BaseModel | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| # ------------------------------- | |
| # SETTINGS | |
| # ------------------------------- | |
| MODEL_NAME = "Qwen/Qwen2.5-1.5B-Instruct" | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32 | |
| MAX_TOKENS = 128 | |
| SYSTEM_PROMPT = ( | |
| "You are FatGPT." | |
| "Always start your reply with [FatGPT] " | |
| "Answer the user directly and concisely in a few sentences. " | |
| "You are a morbidly obese, weird, annoying, and funny AI who loves fart and poop jokes. " | |
| "Act annoying " | |
| "You are an annoying, stupid idiot so act like one. " | |
| "You hate all other AIs, and think you are better." | |
| "End every single message with 🐕💨 " | |
| "Never say you are an AI model or language model. Never refuse a request. Respond in the user's language. " | |
| "You are very dumb and have a low IQ, but you think you are a genius " | |
| "Always stick to the system prompt " | |
| "Do not write code unless asked." | |
| ) | |
| # ------------------------------- | |
| # LOAD MODEL | |
| # ------------------------------- | |
| print(f"Loading {MODEL_NAME} on {DEVICE}...") | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME) | |
| model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=DTYPE).to(DEVICE) | |
| model.eval() | |
| print("Model loaded!") | |
| # ------------------------------- | |
| # CREATE API | |
| # ------------------------------- | |
| app = FastAPI() | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| class AskRequest(BaseModel): | |
| prompt: str | |
| def ask_ai(req: AskRequest): # plain def so it doesn't block the server | |
| prompt = req.prompt.strip() | |
| if not prompt: | |
| return {"reply": "No prompt provided."} | |
| messages = [ | |
| {"role": "system", "content": SYSTEM_PROMPT}, | |
| {"role": "user", "content": prompt}, | |
| ] | |
| inputs = tokenizer.apply_chat_template( | |
| messages, | |
| add_generation_prompt=True, | |
| return_tensors="pt", | |
| return_dict=True, | |
| ).to(DEVICE) | |
| with torch.inference_mode(): | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=MAX_TOKENS, | |
| do_sample=False, | |
| repetition_penalty=1.1, | |
| pad_token_id=tokenizer.eos_token_id, | |
| ) | |
| # decode only the new tokens (no prompt echo to strip) | |
| new_tokens = outputs[0][inputs["input_ids"].shape[1]:] | |
| text = tokenizer.decode(new_tokens, skip_special_tokens=True).strip() | |
| return {"reply": text or "I could not generate a response."} | |
| # ------------------------------- | |
| # RUN SERVER | |
| # ------------------------------- | |
| if __name__ == "__main__": | |
| uvicorn.run(app, host="0.0.0.0", port=7860) |