File size: 1,959 Bytes
f3d004f 4ff5dc7 f3d004f 4ff5dc7 f3d004f 4ff5dc7 f3d004f 4ff5dc7 f3d004f 4ff5dc7 f3d004f 4ff5dc7 f3d004f 4ff5dc7 f3d004f 4ff5dc7 f3d004f 4ff5dc7 f3d004f 4ff5dc7 f3d004f |
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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
# === CONFIG ===
MODEL_NAME = "microsoft/phi-2" # For Free Tier. Use better models for Pro/Ultimate.
# === INIT ===
app = FastAPI(title="FreeAI - Billy (Phi-2)")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
)
model.eval()
# === CORS (so frontend & PWA can call it) ===
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # Replace with your frontend domain in production
allow_methods=["*"],
allow_headers=["*"],
)
# === Request Schema ===
class ChatRequest(BaseModel):
model: str = "free"
messages: list
# === Helper: Format as Chat History ===
def format_prompt(messages):
prompt = "You are Billy, a helpful and friendly assistant.\n\n"
for msg in messages:
role = msg["role"]
content = msg["content"]
if role == "user":
prompt += f"User: {content}\n"
elif role == "assistant":
prompt += f"Billy: {content}\n"
prompt += "Billy:"
return prompt
# === /chat Endpoint (Main Chat API) ===
@app.post("/chat")
async def chat(req: ChatRequest):
prompt = format_prompt(req.messages)
inputs = tokenizer(prompt, return_tensors="pt", truncation=True).to(model.device)
with torch.no_grad():
output = model.generate(
**inputs,
max_new_tokens=150, # Limit for Free Tier
do_sample=True,
temperature=0.7,
top_p=0.9,
pad_token_id=tokenizer.eos_token_id
)
decoded = tokenizer.decode(output[0], skip_special_tokens=True)
response = decoded.split("Billy:")[-1].strip()
return {"message": {"content": response}} |