more initial commit

stick to what I know for now, I guess

Update requirements.txt

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+* text eol=ls

app.py

@@ -0,0 +1,207 @@
+import json
+import time
+import threading
+import gradio as gr
+from spaces import GPU
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from flask import Flask, request, jsonify
+
+gpu = lambda: GPU(duration=120)
+
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16,
+    bnb_4bit_use_double_quant=True,
+)
+
+MODEL_ID = "meta-llama/Llama-3.2-3B-Instruct"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    quantization_config=quantization_config,
+    device_map="auto",
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+)
+
+
+@gpu
+def inference(messages: list, temperature: float, max_tokens: int, top_p: float) -> str:
+    input_ids = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        return_tensors="pt",
+    ).to(model.device)
+
+    do_sample = temperature > 0.0
+    generation_kwargs = {
+        "input_ids": input_ids,
+        "max_new_tokens": max_tokens,
+        "do_sample": do_sample,
+        "pad_token_id": tokenizer.eos_token_id,
+        "eos_token_id": tokenizer.eos_token_id,
+    }
+    if do_sample:
+        generation_kwargs["temperature"] = temperature
+        generation_kwargs["top_p"] = top_p
+
+    with torch.no_grad():
+        output_ids = model.generate(**generation_kwargs)
+
+    new_tokens = output_ids[0][input_ids.shape[1] :]
+    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
+
+
+def run_inference_safe(messages, temperature, max_tokens, top_p):
+    try:
+        return inference(messages, temperature, max_tokens, top_p), None
+    except Exception as e:
+        return None, str(e)
+
+
+def gradio_inference(payload_json: str) -> str:
+    try:
+        payload = json.loads(payload_json)
+    except json.JSONDecodeError as e:
+        return json.dumps({"error": f"Invalid JSON: {e}"})
+    content, err = run_inference_safe(
+        payload.get("messages", []),
+        float(payload.get("temperature", 0.7)),
+        int(payload.get("max_tokens", 1024)),
+        float(payload.get("top_p", 1.0)),
+    )
+    if err:
+        return json.dumps({"error": err})
+    return json.dumps({"content": content})
+
+
+def make_ollama_response(model_name: str, content: str) -> dict:
+    return {
+        "model": model_name,
+        "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+        "message": {
+            "role": "assistant",
+            "content": content,
+        },
+        "done": True,
+    }
+
+
+flask_app = Flask(__name__)
+
+
+@flask_app.route("/api/chat", methods=["POST"])
+def ollama_chat():
+    body = request.get_json(force=True, silent=True) or {}
+    if body.get("stream", False):
+        return jsonify({"error": "Streaming is not supported."}), 400
+
+    messages = body.get("messages", [])
+    model_name = body.get("model", "llama")
+    options = body.get("options", {})
+    temperature = float(options.get("temperature", body.get("temperature", 0.7)))
+    max_tokens = int(options.get("num_predict", body.get("num_predict", 1024)))
+    top_p = float(options.get("top_p", body.get("top_p", 1.0)))
+
+    content, err = run_inference_safe(messages, temperature, max_tokens, top_p)
+    if err:
+        return jsonify({"error": err}), 500
+    return jsonify(make_ollama_response(model_name, content))
+
+
+@flask_app.route("/api/generate", methods=["POST"])
+def ollama_generate():
+    body = request.get_json(force=True, silent=True) or {}
+    if body.get("stream", False):
+        return jsonify({"error": "Streaming is not supported."}), 400
+
+    prompt = body.get("prompt", "")
+    model_name = body.get("model", "llama")
+    options = body.get("options", {})
+    temperature = float(options.get("temperature", 0.7))
+    max_tokens = int(options.get("num_predict", 1024))
+    top_p = float(options.get("top_p", 1.0))
+
+    messages = [{"role": "user", "content": prompt}]
+    content, err = run_inference_safe(messages, temperature, max_tokens, top_p)
+    if err:
+        return jsonify({"error": err}), 500
+    return jsonify(make_ollama_response(model_name, content))
+
+
+@flask_app.route("/api/tags", methods=["GET"])
+def ollama_tags():
+    return jsonify(
+        {
+            "models": [
+                {
+                    "name": "llama",
+                    "model": "llama",
+                    "modified_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
+                    "size": 0,
+                    "digest": "local",
+                    "details": {
+                        "format": "4bit-nf4",
+                        "family": "llama",
+                        "parameter_size": "unknown",
+                        "quantization_level": "Q4_NF4",
+                    },
+                }
+            ]
+        }
+    )
+
+
+@flask_app.route("/v1/models", methods=["GET"])
+def openai_models():
+    return jsonify(
+        {
+            "object": "list",
+            "data": [
+                {
+                    "id": "llama",
+                    "object": "model",
+                    "created": int(time.time()),
+                    "owned_by": "local",
+                }
+            ],
+        }
+    )
+
+
+@flask_app.route("/health", methods=["GET"])
+def health():
+    return jsonify({"status": "ok"})
+
+
+def start_flask():
+    flask_app.run(host="0.0.0.0", port=11434, use_reloader=False)
+
+
+flask_thread = threading.Thread(target=start_flask, daemon=True)
+flask_thread.start()
+
+
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            payload_input = gr.Textbox(
+                label="Request payload (JSON)",
+                placeholder='{"messages": [{"role": "user", "content": "Hello!"}], "max_tokens": 256}',
+                lines=6,
+            )
+            submit_btn = gr.Button("Run inference", variant="primary")
+        with gr.Column():
+            output_box = gr.Textbox(label="Response", lines=6)
+
+    submit_btn.click(
+        fn=gradio_inference,
+        inputs=payload_input,
+        outputs=output_box,
+        api_name="predict",
+    )
+
+demo.launch()