Update app.py

app.py

@@ -1,7 +1,52 @@
+import spaces
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 
-def greet(name):
-    return "Hello " + name + "!!"
+model_id = "PhysicsWallahAI/Aryabhata-1.0"
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id)
+
+stop_strings = [
+    "<|im_end|>",
+    "<|end|>",
+    "<im_start|>",
+    "⁠```python\n",
+    "⁠<|im_start|>",
+    "]}}]}}]",
+]
+
+
+def strip_bad_tokens(s, stop_strings):
+    for suffix in stop_strings:
+        if s.endswith(suffix):
+            return s[: -len(suffix)]
+    return s
+
+
+generation_config = GenerationConfig(max_new_tokens=4096, stop_strings=stop_strings)
+
+
+@spaces.GPU
+def greet(prompt: str):
+    messages = [
+        {
+            "role": "system",
+            "content": "Think step-by-step; put only the final answer inside \\boxed{}.",
+        },
+        {"role": "user", "content": prompt},
+    ]
+    text = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    inputs = tokenizer([text], return_tensors="pt")
+    outputs = model.generate(
+        **inputs, generation_config=generation_config, tokenizer=tokenizer
+    )
+    return strip_bad_tokens(
+        tokenizer.decode(outputs[0], skip_special_tokens=True), stop_strings
+    )
+
+
+demo = gr.Interface(fn=greet, inputs="text", outputs="text", title="Aryabhatta Demo")
 demo.launch()