first app version

.beamignore

@@ -0,0 +1,28 @@
+# Generated by Beam SDK
+.beamignore
+pyproject.toml
+.git
+.idea
+.python-version
+.vscode
+.venv
+venv
+__pycache__
+.DS_Store
+.config
+drive/MyDrive
+.coverage
+.pytest_cache
+.ipynb
+.ruff_cache
+.dockerignore
+.ipynb_checkpoints
+.env.local
+.envrc
+**/__pycache__/
+**/.pytest_cache/
+**/node_modules/
+**/.venv/
+*.pyc
+.next/
+.circleci

.gitattributes

@@ -1,35 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz 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

.gitignore

@@ -0,0 +1,8 @@
+.DS_Store
+__pycache__
+__downloads__
+.env
+.venv/
+.vscode/launch.json
+.secrets
+.idea/

README.md

@@ -4,6 +4,7 @@ emoji: 🚀
 colorFrom: indigo
 colorTo: pink
 sdk: static
+app_file: README.md
 pinned: false
 ---
 

app.py

@@ -0,0 +1,124 @@
+from threading import Thread
+
+import torch
+from beam import Image, Volume, GpuType, asgi
+from fastapi import FastAPI
+from fastapi.responses import StreamingResponse
+from transformers import (
+    AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer,
+    PreTrainedTokenizerFast, PreTrainedModel, StoppingCriteriaList
+)
+
+from utils import MaxPostsStoppingCriteria, Body, fallback
+
+SETTINGS = {
+    "model_name": "Error410/JVCGPT-Medium",
+    "beam_volume_path": "./cached_models",
+}
+
+# @see https://huggingface.co/docs/transformers/generation_strategies#customize-text-generation
+DEFAULTS = {
+    "max_length": 2048,  # 512
+    "temperature": 0.9,  # 1
+    "top_p": 1,  # 0.95
+    "top_k": 0,  # 40
+    "repetition_penalty": 1.0,  # 1.0
+    "no_repeat_ngram_size": 0,  # 0
+    "do_sample": True,  # True
+}
+
+
+def load_models():
+    tokenizer = AutoTokenizer.from_pretrained(
+        SETTINGS["model_name"],
+        cache_dir=SETTINGS["beam_volume_path"]
+    )
+    tokenizer.pad_token = tokenizer.eos_token
+    model = AutoModelForCausalLM.from_pretrained(
+        SETTINGS["model_name"],
+        device_map="auto",
+        torch_dtype=torch.float16,
+        cache_dir=SETTINGS["beam_volume_path"],
+    )
+    return model, tokenizer
+
+
+def stream(model: PreTrainedModel, tokenizer: PreTrainedTokenizerFast, body: Body):
+    generate_args = {
+        "max_length": fallback(body.max_length, DEFAULTS["max_length"]),
+        "temperature": fallback(body.temperature, DEFAULTS["temperature"]),
+        "top_p": fallback(body.top_p, DEFAULTS["top_p"]),
+        "top_k": fallback(body.top_k, DEFAULTS["top_k"]),
+        "repetition_penalty": fallback(body.repetition_penalty, DEFAULTS["repetition_penalty"]),
+        "no_repeat_ngram_size": fallback(body.no_repeat_ngram_size, DEFAULTS["no_repeat_ngram_size"]),
+        "do_sample": fallback(body.do_sample, DEFAULTS["do_sample"]),
+        "use_cache": True,
+        "eos_token_id": tokenizer.eos_token_id,
+        "pad_token_id": tokenizer.pad_token_id,
+    }
+
+    inputs = tokenizer(body.prompt, return_tensors="pt", padding=True)
+    input_ids = inputs["input_ids"].to("cuda")
+    attention_mask = inputs["attention_mask"].to("cuda")
+
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False, timeout=240)
+
+    # with torch.no_grad(): # seems to be useless
+    thread = Thread(
+        target=model.generate,
+        kwargs={
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "streamer": streamer,
+            "stopping_criteria": StoppingCriteriaList([MaxPostsStoppingCriteria(tokenizer, body.posts_count)]),
+            **generate_args,
+        }
+    )
+    thread.start()
+
+    for token in streamer:
+        yield token
+    #     if len(token) > 0:
+    #         yield f"DATA {token}"
+    #
+    # yield "EOS"
+
+
+@asgi(
+    on_start=load_models,
+    cpu=2.0,
+    memory="16Gi",
+    gpu=GpuType.A100_40,
+    gpu_count=1,
+    timeout=900,  # Time for loading the model and run the server
+    image=Image(
+        python_version="python3.12",
+        python_packages=[
+            "fastapi",
+            "torch",
+            "transformers",
+            "accelerate",
+            "huggingface_hub[hf-transfer]",
+        ],
+        env_vars=["HF_HUB_ENABLE_HF_TRANSFER=1"],
+    ),
+    volumes=[
+        Volume(
+            name="cached_models",
+            mount_path=SETTINGS["beam_volume_path"],
+        )
+    ],
+)
+def server(context):
+    model, tokenizer = context.on_start_value
+    app = FastAPI()
+
+    @app.post("/stream")
+    async def stream_endpoint(body: Body) -> StreamingResponse:
+        return StreamingResponse(
+            stream(model, tokenizer, body),
+            media_type='text/event-stream',
+            headers={"Cache-Control": "no-cache"},
+        )
+
+    return app

index.html

@@ -1,19 +0,0 @@
-<!doctype html>
-<html>
-	<head>
-		<meta charset="utf-8" />
-		<meta name="viewport" content="width=device-width" />
-		<title>My static Space</title>
-		<link rel="stylesheet" href="style.css" />
-	</head>
-	<body>
-		<div class="card">
-			<h1>Welcome to your static Space!</h1>
-			<p>You can modify this app directly by editing <i>index.html</i> in the Files and versions tab.</p>
-			<p>
-				Also don't forget to check the
-				<a href="https://huggingface.co/docs/hub/spaces" target="_blank">Spaces documentation</a>.
-			</p>
-		</div>
-	</body>
-</html>

style.css

@@ -1,28 +0,0 @@
-body {
-	padding: 2rem;
-	font-family: -apple-system, BlinkMacSystemFont, "Arial", sans-serif;
-}
-
-h1 {
-	font-size: 16px;
-	margin-top: 0;
-}
-
-p {
-	color: rgb(107, 114, 128);
-	font-size: 15px;
-	margin-bottom: 10px;
-	margin-top: 5px;
-}
-
-.card {
-	max-width: 620px;
-	margin: 0 auto;
-	padding: 16px;
-	border: 1px solid lightgray;
-	border-radius: 16px;
-}
-
-.card p:last-child {
-	margin-bottom: 0;
-}

utils.py

@@ -0,0 +1,36 @@
+from pydantic import BaseModel
+from transformers import (PreTrainedTokenizerFast, StoppingCriteria)
+
+
+def fallback(value, fallback_value):
+    if value is None:
+        return fallback_value
+    return value
+
+
+class Body(BaseModel):
+    prompt: str
+    posts_count: int
+    max_length: int | None = None
+    temperature: float | None = None
+    top_p: float | None = None
+    top_k: float | None = None
+    repetition_penalty: float | None = None
+    no_repeat_ngram_size: float | None = None
+    do_sample: bool | None = None
+
+
+class MaxPostsStoppingCriteria(StoppingCriteria):
+    def __init__(self, tokenizer: PreTrainedTokenizerFast, posts_count: int):
+        self.end_of_post_token_id = tokenizer.encode("<|end_of_post|>", add_special_tokens=False)
+        self.posts_count = posts_count
+        self.counter = 0
+
+    def __call__(self, input_ids, scores, **kwargs):
+        # Check if the last token matches the <|end_of_post|> token ID
+        for sequence in input_ids:
+            if sequence[-len(self.end_of_post_token_id):].tolist() == self.end_of_post_token_id:
+                self.counter += 1
+                if self.counter >= self.posts_count:
+                    return True
+        return False