Update app.py

app.py

@@ -1,36 +1,36 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+import torch
 
-model_id = "tiiuae/falcon-rw-1b"  # lightweight model that works on CPU
-
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id)
+model_name = "ajibawa-2023/Young-Children-Storyteller-Mistral-7B"
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4"
+)
 
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    quantization_config=quant_config,
+    device_map="auto",
+    torch_dtype=torch.float16,
+    trust_remote_code=True
+)
 
-def generate_story(character_1, character_2, theme, setting, scenario):
-    prompt = (
-        f"Characters: {character_1}, {character_2}\n"
-        f"Theme: {theme}\n"
-        f"Setting: {setting}\n"
-        f"Scenario: {scenario}\n"
-        f"Write a story suitable for children aged 6 to 12:"
+def generate_story(prompt, max_length=400, temperature=0.7, top_p=0.9):
+    formatted_prompt = f"### Instruction:\nCreate a story for young children about: {prompt}\n\n### Response:\n"
+    inputs = tokenizer.encode(formatted_prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(
+        inputs,
+        max_length=max_length,
+        temperature=temperature,
+        top_p=top_p,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id,
+        repetition_penalty=1.1
     )
-    output = generator(prompt, max_new_tokens=300, temperature=0.8)[0]["generated_text"]
-    return output[len(prompt):].strip()
-
-demo = gr.Interface(
-    fn=generate_story,
-    inputs=[
-        gr.Textbox(label="Character 1"),
-        gr.Textbox(label="Character 2"),
-        gr.Textbox(label="Theme"),
-        gr.Textbox(label="Setting"),
-        gr.Textbox(label="Scenario")
-    ],
-    outputs=gr.Textbox(label="📖 Generated Story"),
-    title="Genieverse Lite Story Generator",
-    description="Enter characters, theme, and scenario to create a children's story!"
-)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True).split("### Response:")[-1].strip()
 
-demo.launch(server_name="0.0.0.0", server_port=7860)
+gr.Interface(fn=generate_story, inputs="text", outputs="text").launch()