|
|
from flask import Flask, request, jsonify, render_template_string |
|
|
from flask_cors import CORS |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
import time |
|
|
import os |
|
|
|
|
|
|
|
|
torch.set_num_threads(1) |
|
|
torch.set_num_interop_threads(1) |
|
|
os.environ['OMP_NUM_THREADS'] = '1' |
|
|
os.environ['MKL_NUM_THREADS'] = '1' |
|
|
|
|
|
|
|
|
class CharTokenizer: |
|
|
def __init__(self): |
|
|
chars = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ " |
|
|
chars += "0123456789.,!?¿áéíóúñÁÉÍÓÚÑ" |
|
|
self.char_to_idx = {c: i+1 for i, c in enumerate(chars)} |
|
|
self.idx_to_char = {i+1: c for i, c in enumerate(chars)} |
|
|
self.vocab_size = len(self.char_to_idx) + 1 |
|
|
|
|
|
def encode(self, text, max_len=100): |
|
|
indices = [self.char_to_idx.get(c, 0) for c in text[:max_len]] |
|
|
indices += [0] * (max_len - len(indices)) |
|
|
return torch.tensor(indices, dtype=torch.long) |
|
|
|
|
|
class AtacamaWeatherOracle(nn.Module): |
|
|
def __init__(self, vocab_size=100, embed_dim=16, hidden_dim=32): |
|
|
super().__init__() |
|
|
self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0) |
|
|
self.lstm = nn.LSTM(embed_dim, hidden_dim, batch_first=True) |
|
|
self.classifier = nn.Linear(hidden_dim, 2) |
|
|
|
|
|
def forward(self, x): |
|
|
embedded = self.embedding(x) |
|
|
_, (hidden, _) = self.lstm(embedded) |
|
|
logits = self.classifier(hidden.squeeze(0)) |
|
|
return logits |
|
|
|
|
|
|
|
|
app = Flask(__name__) |
|
|
CORS(app) |
|
|
|
|
|
|
|
|
print("Loading Atacama Weather Oracle...") |
|
|
load_start = time.time() |
|
|
tokenizer = CharTokenizer() |
|
|
model = AtacamaWeatherOracle(vocab_size=tokenizer.vocab_size) |
|
|
|
|
|
checkpoint = torch.load('atacama_weather_oracle.pth', map_location='cpu') |
|
|
model.load_state_dict(checkpoint['model_state_dict']) |
|
|
model.eval() |
|
|
load_time = time.time() - load_start |
|
|
print(f"✅ Oracle loaded and ready! (took {load_time:.3f}s)") |
|
|
|
|
|
|
|
|
HTML_TEMPLATE = """ |
|
|
<!DOCTYPE html> |
|
|
<html lang="en"> |
|
|
<head> |
|
|
<meta charset="UTF-8"> |
|
|
<meta name="viewport" content="width=device-width, initial-scale=1.0"> |
|
|
<title>Is It Raining in Atacama?</title> |
|
|
<style> |
|
|
* { |
|
|
margin: 0; |
|
|
padding: 0; |
|
|
box-sizing: border-box; |
|
|
} |
|
|
body { |
|
|
font-family: 'Courier New', monospace; |
|
|
max-width: 700px; |
|
|
margin: 0 auto; |
|
|
padding: 40px 20px; |
|
|
background: #fafafa; |
|
|
color: #1a1a1a; |
|
|
line-height: 1.6; |
|
|
} |
|
|
.container { |
|
|
background: white; |
|
|
padding: 40px; |
|
|
border: 1px solid #e0e0e0; |
|
|
} |
|
|
h1 { |
|
|
font-size: 1.5em; |
|
|
font-weight: normal; |
|
|
margin-bottom: 8px; |
|
|
letter-spacing: -0.5px; |
|
|
} |
|
|
.subtitle { |
|
|
font-size: 0.85em; |
|
|
color: #666; |
|
|
margin-bottom: 30px; |
|
|
font-family: -apple-system, sans-serif; |
|
|
} |
|
|
.stats { |
|
|
display: inline-block; |
|
|
background: #f5f5f5; |
|
|
padding: 2px 8px; |
|
|
margin: 0 4px; |
|
|
font-size: 0.8em; |
|
|
border-radius: 2px; |
|
|
} |
|
|
input[type="text"] { |
|
|
width: 100%; |
|
|
padding: 12px; |
|
|
font-size: 15px; |
|
|
font-family: -apple-system, sans-serif; |
|
|
border: 1px solid #d0d0d0; |
|
|
margin-bottom: 12px; |
|
|
background: #fafafa; |
|
|
} |
|
|
input[type="text"]:focus { |
|
|
outline: none; |
|
|
border-color: #1a1a1a; |
|
|
background: white; |
|
|
} |
|
|
button { |
|
|
width: 100%; |
|
|
padding: 12px; |
|
|
font-size: 15px; |
|
|
font-family: 'Courier New', monospace; |
|
|
background: #1a1a1a; |
|
|
color: white; |
|
|
border: none; |
|
|
cursor: pointer; |
|
|
transition: background 0.2s; |
|
|
} |
|
|
button:hover { |
|
|
background: #333; |
|
|
} |
|
|
#result { |
|
|
margin-top: 30px; |
|
|
padding: 20px; |
|
|
background: #f9f9f9; |
|
|
border-left: 3px solid #1a1a1a; |
|
|
display: none; |
|
|
font-family: -apple-system, sans-serif; |
|
|
} |
|
|
.answer { |
|
|
font-size: 2em; |
|
|
font-weight: 300; |
|
|
margin-bottom: 8px; |
|
|
font-family: 'Courier New', monospace; |
|
|
} |
|
|
.confidence { |
|
|
font-size: 0.9em; |
|
|
color: #666; |
|
|
} |
|
|
.footer { |
|
|
margin-top: 40px; |
|
|
padding-top: 20px; |
|
|
border-top: 1px solid #e0e0e0; |
|
|
font-size: 0.8em; |
|
|
color: #999; |
|
|
font-family: -apple-system, sans-serif; |
|
|
} |
|
|
.emoji { |
|
|
font-size: 2em; |
|
|
margin-bottom: 10px; |
|
|
} |
|
|
.timing { |
|
|
margin-top: 10px; |
|
|
font-size: 0.75em; |
|
|
color: #aaa; |
|
|
font-family: 'Courier New', monospace; |
|
|
} |
|
|
</style> |
|
|
</head> |
|
|
<body> |
|
|
<div class="container"> |
|
|
<h1>atacama</h1> |
|
|
<p class="subtitle"> |
|
|
An ultra-small language model |
|
|
<span class="stats">7,762 parameters</span> |
|
|
<span class="stats">30KB</span> |
|
|
<span class="stats">99.9% certain</span> |
|
|
</p> |
|
|
|
|
|
<input type="text" id="question" placeholder="is it raining in atacama?" |
|
|
value="is it raining in atacama?"> |
|
|
<button onclick="askOracle()">ask</button> |
|
|
|
|
|
<div id="result"></div> |
|
|
|
|
|
<div class="footer"> |
|
|
trained on 50+ years of atacama desert weather data<br> |
|
|
last recorded rainfall: march 2015 |
|
|
</div> |
|
|
</div> |
|
|
|
|
|
<script> |
|
|
async function askOracle() { |
|
|
const question = document.getElementById('question').value; |
|
|
const resultDiv = document.getElementById('result'); |
|
|
|
|
|
resultDiv.style.display = 'block'; |
|
|
resultDiv.innerHTML = '<p>Consulting the oracle...</p>'; |
|
|
|
|
|
const startTime = performance.now(); |
|
|
|
|
|
try { |
|
|
const response = await fetch('/ask', { |
|
|
method: 'POST', |
|
|
headers: {'Content-Type': 'application/json'}, |
|
|
body: JSON.stringify({question: question}) |
|
|
}); |
|
|
|
|
|
const endTime = performance.now(); |
|
|
const totalTime = ((endTime - startTime) / 1000).toFixed(2); |
|
|
|
|
|
const data = await response.json(); |
|
|
|
|
|
const emoji = data.prob_no_rain > 0.999 ? '☀️' : '🌤️'; |
|
|
|
|
|
resultDiv.innerHTML = ` |
|
|
<div class="emoji">${emoji}</div> |
|
|
<div class="answer">${data.answer}</div> |
|
|
<div class="confidence">${data.confidence}</div> |
|
|
<div class="confidence" style="margin-top: 10px; font-size: 0.9em;"> |
|
|
No rain: ${(data.prob_no_rain * 100).toFixed(2)}% | |
|
|
Rain: ${(data.prob_rain * 100).toFixed(2)}% |
|
|
</div> |
|
|
<div class="timing"> |
|
|
⏱️ total: ${totalTime}s | server inference: ${data.inference_ms}ms |
|
|
</div> |
|
|
`; |
|
|
} catch (error) { |
|
|
resultDiv.innerHTML = '<p>Error: Could not reach the oracle</p>'; |
|
|
} |
|
|
} |
|
|
|
|
|
// Allow Enter key to submit |
|
|
document.getElementById('question').addEventListener('keypress', function(e) { |
|
|
if (e.key === 'Enter') askOracle(); |
|
|
}); |
|
|
</script> |
|
|
</body> |
|
|
</html> |
|
|
""" |
|
|
|
|
|
@app.route('/') |
|
|
def home(): |
|
|
return render_template_string(HTML_TEMPLATE) |
|
|
|
|
|
@app.route('/ask', methods=['POST']) |
|
|
def ask(): |
|
|
request_start = time.time() |
|
|
|
|
|
data = request.json |
|
|
question = data.get('question', '') |
|
|
|
|
|
|
|
|
t0 = time.time() |
|
|
tokens = tokenizer.encode(question).unsqueeze(0) |
|
|
t1 = time.time() |
|
|
|
|
|
with torch.no_grad(): |
|
|
logits = model(tokens) |
|
|
t2 = time.time() |
|
|
|
|
|
probs = torch.softmax(logits, dim=1)[0] |
|
|
t3 = time.time() |
|
|
|
|
|
prob_no_rain = probs[0].item() |
|
|
prob_rain = probs[1].item() |
|
|
t4 = time.time() |
|
|
|
|
|
if prob_no_rain > 0.999: |
|
|
answer = "No." |
|
|
confidence = "Absolute certainty" |
|
|
elif prob_no_rain > 0.99: |
|
|
answer = "No. (But I admire your optimism)" |
|
|
confidence = "Very high confidence" |
|
|
elif prob_no_rain > 0.9: |
|
|
answer = "Almost certainly not." |
|
|
confidence = "High confidence" |
|
|
else: |
|
|
answer = "Historically unprecedented... but no." |
|
|
confidence = "Moderate confidence" |
|
|
|
|
|
total_time = time.time() - request_start |
|
|
|
|
|
|
|
|
print(f"TIMING: tokenize={((t1-t0)*1000):.1f}ms, model={((t2-t1)*1000):.1f}ms, softmax={((t3-t2)*1000):.1f}ms, extract={((t4-t3)*1000):.1f}ms, total={total_time*1000:.1f}ms") |
|
|
|
|
|
return jsonify({ |
|
|
'answer': answer, |
|
|
'confidence': confidence, |
|
|
'prob_no_rain': prob_no_rain, |
|
|
'prob_rain': prob_rain, |
|
|
'inference_ms': f"{total_time*1000:.1f}", |
|
|
'debug': f"tok={((t1-t0)*1000):.0f}ms model={((t2-t1)*1000):.0f}ms soft={((t3-t2)*1000):.0f}ms" |
|
|
}) |
|
|
|
|
|
@app.route('/health') |
|
|
def health(): |
|
|
"""Health check endpoint - also useful for keeping the container warm""" |
|
|
return jsonify({'status': 'ok', 'model': 'loaded'}) |
|
|
|
|
|
if __name__ == '__main__': |
|
|
import os |
|
|
port = int(os.environ.get('PORT', 5000)) |
|
|
app.run(host='0.0.0.0', port=port) |
|
|
|
