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