import asyncio
from transformers import AutoModelForCausalLM, AutoTokenizer
import websockets

# 모델 초기화
model_name = "Qwen/Qwen2.5-Coder-32B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", torch_dtype="auto")

# WebSocket 서버 생성
async def chat_server(websocket, path):
    print("Client connected.")
    while True:
        # 클라이언트로부터 메시지 수신
        input_text = await websocket.recv()
        print(f"User: {input_text}")

        # 모델 출력 생성
        inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
        output = model.generate(**inputs, max_new_tokens=100)
        response = tokenizer.decode(output[0], skip_special_tokens=True)

        # 응답 전송
        await websocket.send(response)
        print(f"Model: {response}")

# WebSocket 서버 시작
async def main():
    async with websockets.serve(chat_server, "localhost", 8765):
        print("Server started on ws://localhost:8765")
        await asyncio.Future()  # 서버를 종료하지 않도록 유지

asyncio.run(main())

import asyncio
import websockets

async def chat():
    uri = "ws://localhost:8765"
    async with websockets.connect(uri) as websocket:
        while True:
            user_input = input("You: ")
            await websocket.send(user_input)
            response = await websocket.recv() 
            print(f"AI: {response}")

asyncio.run(chat())