arthu1's picture
Update app.py
64d3bbc verified
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
# Load your model
print("πŸš€ Loading Starlight Mini...")
MODEL_NAME = "arthu1/starlight-mini"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
torch_dtype=torch.float16,
device_map="auto",
low_cpu_mem_usage=True
)
print("βœ… Model loaded!")
def generate_text(
prompt: str,
max_tokens: int = 100,
temperature: float = 0.7,
top_p: float = 0.9
) -> str:
"""Generate text from prompt"""
if not prompt:
return "Error: Please provide a prompt"
try:
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=min(max_tokens, 500),
temperature=temperature,
top_p=top_p,
do_sample=True,
pad_token_id=tokenizer.eos_token_id
)
result = tokenizer.decode(outputs[0], skip_special_tokens=True)
return result
except Exception as e:
return f"Error: {str(e)}"
# Create Gradio interface
with gr.Blocks(title="Starlight Mini API") as demo:
gr.Markdown("# 🌟 Starlight Mini API\n**Free AI Text Generation - 8B Parameters**")
with gr.Tab("πŸ’¬ Chat"):
with gr.Row():
with gr.Column():
prompt_input = gr.Textbox(label="Prompt", placeholder="Enter your prompt...", lines=5)
with gr.Accordion("Settings", open=False):
max_tokens_slider = gr.Slider(10, 500, 100, step=10, label="Max Tokens")
temperature_slider = gr.Slider(0.1, 2.0, 0.7, step=0.1, label="Temperature")
top_p_slider = gr.Slider(0.1, 1.0, 0.9, step=0.05, label="Top P")
generate_btn = gr.Button("✨ Generate", variant="primary")
with gr.Column():
output_text = gr.Textbox(label="Output", lines=10)
gr.Examples(
[["Once upon a time"], ["Explain AI:"], ["Code a function:"]],
inputs=[prompt_input]
)
with gr.Tab("πŸ“‘ API Docs"):
gr.Markdown("""
## API Usage
**Python:**
```python
from gradio_client import Client
client = Client("arthu1/Exquisite-Starlight")
result = client.predict("Hello!", 100, 0.7, 0.9, api_name="/predict")
```
**cURL:**
```bash
curl -X POST https://arthu1-exquisite-starlight.hf.space/api/predict \\
-H "Content-Type: application/json" \\
-d '{"data": ["Hello!", 100, 0.7, 0.9]}'
```
""")
generate_btn.click(
fn=generate_text,
inputs=[prompt_input, max_tokens_slider, temperature_slider, top_p_slider],
outputs=output_text
)
demo.launch(show_api=True, share=True)