Create app.py

app.py

@@ -0,0 +1,68 @@
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32
+)
+if torch.cuda.is_available():
+    model.to("cuda")
+model.eval()
+
+def generate_text(prompt, max_new_tokens=100, temperature=0.7, top_k=50):
+    if not prompt:
+        return "Please enter a prompt."
+
+    messages = [{"role": "user", "content": prompt}]
+    encoded = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        return_tensors="pt",
+        padding=True,
+        return_attention_mask=True,
+    )
+
+    input_ids = encoded["input_ids"]
+    attention_mask = encoded["attention_mask"]
+
+    if torch.cuda.is_available():
+        input_ids = input_ids.to("cuda")
+        attention_mask = attention_mask.to("cuda")
+
+    output_ids = model.generate(
+        input_ids,
+        attention_mask=attention_mask,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,
+        temperature=temperature,
+        top_k=top_k,
+        pad_token_id=tokenizer.eos_token_id
+    )
+
+    response = tokenizer.decode(output_ids[0][input_ids.shape[-1]:], skip_special_tokens=True)
+    return response
+
+# Gradio interface
+demo = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.Textbox(label="Prompt"),
+        gr.Slider(minimum=10, maximum=500, value=100, label="Max New Tokens"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.05, label="Temperature"),
+        gr.Slider(minimum=0, maximum=100, value=50, step=1, label="Top K")
+    ],
+    outputs=gr.Textbox(label="Generated Text"),
+    title="TinyLlama Gradio API",
+    description="Use this via UI or API via `/run/predict`"
+)
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)