File size: 1,170 Bytes
93009f8 52642dd 93009f8 52642dd 93009f8 52642dd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 |
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
# Load the model
model_id = "mistralai/Mistral-7B-Instruct-v0.2" # or another instruct-tuned model
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype="auto")
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
# AI Function
def generate_json(prompt):
system_prompt = "You are an AI that generates valid and clean JSON objects based on descriptions."
input_prompt = f"{system_prompt}\n\nDescription:\n{prompt}\n\nJSON:\n"
output = generator(input_prompt, max_new_tokens=256, do_sample=False)[0]["generated_text"]
# Remove original prompt from generated text
json_part = output.replace(input_prompt, "").strip()
return json_part
# Gradio UI
gr.Interface(
fn=generate_json,
inputs=gr.Textbox(label="Describe the JSON you want", lines=4),
outputs=gr.Textbox(label="Generated JSON", lines=20),
title="🧠 JSON Generator with LLM",
description="Enter a plain language description and get a structured JSON output."
).launch()
|