Update app.py

app.py

@@ -1,29 +1,37 @@
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse, StreamingResponse
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
 from threading import Thread
 from collections import defaultdict
-import torch
 
-app = FastAPI(title="Mariza Koller 1.5B - API com Memória 😈")
+app = FastAPI(title="Mariza Koller 1.5B - CPU Free 4bit 🔥")
 
-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)
+print("🔥 Carregando Qwen2-1.5B em 4-bit na CPU (agora SIM funciona no HF free!)")
+
+# Configuração 4-bit que roda na CPU do Spaces
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype="float16",
+    bnb_4bit_use_double_quant=True,
+)
+
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2-1.5B-Instruct", trust Tom's_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
-    "Qwen/Qwen2-1.5B-Instruct",
+    "Qwen/Qwen2-1.5B-Instruct2",
     device_map="cpu",
-    load_in_8bit=True,
-    torch_dtype=torch.float16,
-    trust_remote_code=True
+    quantization_config=quantization_config,
+    trust_remote_code=True,
+    low_cpu_mem_usage=True
 )
 
-# Cache de conversa em memória: {user_id: lista de mensagens}
+# Cache de conversa
 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"}
+    return {"message": "Mariza 1.5B 4-bit rodando quente na CPU free, chefe! 😈"}
 
 @app.post("/chat")
 async def chat(request: Request):
@@ -35,9 +43,9 @@ async def chat(request: Request):
     stream = data.get("stream", False)
 
     if not prompt:
-        return JSONResponse({"error": "prompt vazio, seu safado"})
+        return JSONResponse({"error": "manda prompt direito, safado"})
 
-    # Monta histórico no formato do Qwen2
+    # Monta histórico
     messages = history_db[user_id]
     full_prompt = ""
     for role, content in messages:
@@ -60,16 +68,10 @@ async def chat(request: Request):
         }
         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")
-
+        return StreamingResponse(streamer, media_type="text/event-stream")
     else:
         outputs = model.generate(
-            input_ids=inputs.input_ids,
-            attention_mask=inputs.attention_mask,
+            **inputs,
             max_new_tokens=max_tokens,
             temperature=temperature,
             do_sample=True,
@@ -79,12 +81,13 @@ async def chat(request: Request):
         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:
+        # Limpa histórico se ficar grande
+        while sum(len(tokenizer.encode(c[1])) for c in messages) > MAX_CONTEXT_TOKENS:
             messages.pop(0)
 
-        return JSONResponse({"response": resposta, "user_id": user_id})
+        return JSONResponse({"response": resposta})
+
+print("✅ Modelo carregado! Mariza tá pronta pra dominar o WhatsApp 😏")