Update app.py

app.py

@@ -27,6 +27,53 @@ async def test_endpoint(message: dict):
     return response
 
 
+@app.post("/chat/")
+async def chat_endpoint(message: dict):
+    if "text" not in message:
+        raise HTTPException(status_code=400, detail="Missing 'text' in request body")
+    
+    chat_message = message["text"]
+    response_text = generate_chat_response(chat_message)
+    
+    return {"response": response_text}
+
+def generate_chat_response(text: str):
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.bfloat16,
+        low_cpu_mem_usage=True,
+        trust_remote_code=True
+    )
+    
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+
+    conversation = [{"role": "user", "content": text}]
+    input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True,
+                                              return_tensors="pt", return_dict=True).to(model.device)
+    streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
+
+    generate_kwargs = dict(
+        max_length=4096,
+        streamer=streamer,
+        do_sample=True,
+        top_p=0.9,
+        top_k=50,
+        temperature=0.7,
+        repetition_penalty=1.0,
+        eos_token_id=[151329, 151336, 151338],
+    )
+    gen_kwargs = {**input_ids, **generate_kwargs}
+
+    with torch.no_grad():
+        thread = Thread(target=model.generate, kwargs=gen_kwargs)
+        thread.start()
+        buffer = ""
+        for new_text in streamer:
+            buffer += new_text
+    
+    return buffer
+
+
 MODEL_LIST = ["nikravan/glm-4vq"]
 
 HF_TOKEN = os.environ.get("HF_TOKEN", None)