Create app.py

app.py

@@ -0,0 +1,45 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+# Load model and tokenizer
+model = AutoModelForCausalLM.from_pretrained(
+    "FlameF0X/SnowflakeCore-G1-Tiny",
+    trust_remote_code=True,
+    force_download=True,
+    use_safetensors=True,
+)
+tokenizer = AutoTokenizer.from_pretrained(
+    "FlameF0X/SnowflakeCore-G1-Tiny",
+    trust_remote_code=True,
+    force_download=True,
+    use_safetensors=True,
+)
+
+def custom_greedy_generate(prompt, max_length=50):
+    model.eval()
+    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
+    generated = input_ids
+    with torch.no_grad():
+        for _ in range(max_length):
+            outputs = model(input_ids=generated)
+            next_token_logits = outputs["logits"][:, -1, :]
+            next_token_id = torch.argmax(next_token_logits, dim=-1).unsqueeze(-1)
+            generated = torch.cat((generated, next_token_id), dim=1)
+            if next_token_id.item() == tokenizer.eos_token_id:
+                break
+    return tokenizer.decode(generated[0], skip_special_tokens=True)
+
+def gradio_generate(prompt):
+    return custom_greedy_generate(prompt)
+
+iface = gr.Interface(
+    fn=gradio_generate,
+    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt here..."),
+    outputs=gr.Textbox(label="Generated Text"),
+    title="SnowflakeCore-G1-Tiny Text Generation",
+    description="Enter a prompt and generate text using the SnowflakeCore-G1-Tiny model.",
+)
+
+if __name__ == "__main__":
+    iface.launch()