| from fastapi import FastAPI, WebSocket, WebSocketDisconnect |
| from fastapi.responses import HTMLResponse |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline |
|
|
| |
| model_name = "bigcode/starcoder2-7b" |
| tokenizer = AutoTokenizer.from_pretrained(model_name) |
| model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto") |
| pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device_map="auto") |
|
|
| app = FastAPI() |
|
|
| |
| html = """ |
| <!DOCTYPE html> |
| <html lang="en"> |
| <head> |
| <meta charset="UTF-8"> |
| <title>StarCoder2 Docker Terminal</title> |
| <style> |
| body { background-color: #1e1e1e; color: #00ff99; font-family: monospace; margin: 0; height: 100vh; display:flex; justify-content:center; align-items:center; } |
| #terminal { width: 80%; height: 90%; } |
| </style> |
| <script src="https://cdn.jsdelivr.net/npm/xterm@5.3.0/lib/xterm.js"></script> |
| <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/xterm@5.3.0/css/xterm.css" /> |
| </head> |
| <body> |
| <div id="terminal"></div> |
| <script> |
| const term = new Terminal(); |
| term.open(document.getElementById('terminal')); |
| term.write("🐳 StarCoder2 Terminal\\nroot@docker:~$ "); |
| |
| const ws = new WebSocket(`ws://${window.location.host}/ws`); |
| |
| term.onKey(e => { |
| if (e.domEvent.key === "Enter") { |
| const command = term.buffer.active.getLine(term.buffer.active.cursorY).translateToString().replace("root@docker:~$ ","").trim(); |
| ws.send(command); |
| term.write("\\r\\n"); |
| } else { |
| term.write(e.key); |
| } |
| }); |
| |
| ws.onmessage = (event) => { |
| term.write(event.data + "\\r\\nroot@docker:~$ "); |
| }; |
| </script> |
| </body> |
| </html> |
| """ |
|
|
| @app.get("/") |
| async def get_root(): |
| return HTMLResponse(html) |
|
|
| @app.websocket("/ws") |
| async def websocket_endpoint(ws: WebSocket): |
| await ws.accept() |
| history = [] |
| try: |
| while True: |
| user_input = await ws.receive_text() |
| prompt = "".join([f"User: {h[0]}\\nAssistant: {h[1]}\\n" for h in history]) |
| prompt += f"User: {user_input}\\nAssistant:" |
|
|
| output = pipe(prompt, max_new_tokens=256, temperature=0.7, pad_token_id=tokenizer.eos_token_id) |
| reply = output[0]["generated_text"][len(prompt):].strip() |
| history.append((user_input, reply)) |
|
|
| await ws.send_text(reply) |
| except WebSocketDisconnect: |
| print("Client disconnected") |
|
|
