Create app.py

app.py

@@ -0,0 +1,90 @@
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from threading import Thread
+from collections import defaultdict
+import torch
+
+app = FastAPI(title="Mariza Koller 1.5B - API com Memória 😈")
+
+print("🔥 Carregando Qwen2-1.5B-Instruct em int8 na CPU... (aguenta aí 2-3 min na primeira vez)")
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B-Instruct", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    "Qwen/Qwen2-1.5B-Instruct",
+    device_map="cpu",
+    load_in_8bit=True,
+    torch_dtype=torch.float16,
+    trust_remote_code=True
+)
+
+# Cache de conversa em memória: {user_id: lista de mensagens}
+history_db = defaultdict(list)
+MAX_CONTEXT_TOKENS = 3500
+
+@app.get("/")
+async def root():
+    return {"message": "Mariza 1.5B tá viva e quente na CPU, chefe! 😏 manda POST /chat"}
+
+@app.post("/chat")
+async def chat(request: Request):
+    data = await request.json()
+    prompt = data.get("prompt", "").strip()
+    user_id = str(data.get("user_id", "default"))
+    max_tokens = data.get("max_tokens", 512)
+    temperature = data.get("temperature", 0.7)
+    stream = data.get("stream", False)
+
+    if not prompt:
+        return JSONResponse({"error": "prompt vazio, seu safado"})
+
+    # Monta histórico no formato do Qwen2
+    messages = history_db[user_id]
+    full_prompt = ""
+    for role, content in messages:
+        full_prompt += f"<|im_start|>{role}\n{content}<|im_end|>\n"
+    full_prompt += f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
+
+    inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=4096)
+
+    if stream:
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = {
+            "input_ids": inputs.input_ids,
+            "attention_mask": inputs.attention_mask,
+            "max_new_tokens": max_tokens,
+            "temperature": temperature,
+            "do_sample": True,
+            "top_p": 0.9,
+            "repetition_penalty": 1.1,
+            "streamer": streamer
+        }
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
+        thread.start()
+
+        def generate():
+            for new_text in streamer:
+                yield new_text
+        return StreamingResponse(generate(), media_type="text/event-stream")
+
+    else:
+        outputs = model.generate(
+            input_ids=inputs.input_ids,
+            attention_mask=inputs.attention_mask,
+            max_new_tokens=max_tokens,
+            temperature=temperature,
+            do_sample=True,
+            top_p=0.9,
+            repetition_penalty=1.1
+        )
+        resposta = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        resposta = resposta.split("<|im_start|>assistant")[-1].strip()
+
+        # Salva no histórico
+        messages.append(("user", prompt))
+        messages.append(("assistant", resposta))
+
+        # Limpa histórico antigo se passar do limite
+        while sum(len(tokenizer.encode(m[1])) for m in messages) > MAX_CONTEXT_TOKENS:
+            messages.pop(0)
+
+        return JSONResponse({"response": resposta, "user_id": user_id})