| | 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") |
| |
|
