Change model from bloom to bloomz and change prompt (#1)

- Update app.py (7bc648a82986fa4d32ac01272acbf13e465bcdb5)


Co-authored-by: Sarah <[email protected]>

app.py

@@ -1,55 +1,14 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-MODEL_ID = "bigscience/bloom-560m"
-
-# load once globally
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    device_map="auto",   # GPU if available, otherwise CPU
-    trust_remote_code=True
-)
-
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=200,
-    temperature=0.7
-)
-
-def generate_idiom(situation: str):
-    prompt = f"""You are a wise assistant. Given a situation, respond with exactly:
-1. A single Chinese idiom (成语).
-2. Its pinyin.
-3. A short English explanation.
-
-Format:
-Idiom
-Pinyin
-Explanation
-
-Situation: {situation}
-Answer:
-"""
-    response = pipe(prompt)[0]["generated_text"]
-    clean_response = response.split("Answer:")[-1].strip()
-
-    # Try to split into lines
-    lines = [line.strip() for line in clean_response.split("\n") if line.strip()]
-    if len(lines) >= 3:
-        idiom = lines[0]
-        pinyin = lines[1]
-        meaning = " ".join(lines[2:])
-        explanation = f"{pinyin}\n\n{meaning}"
-    else:
-        # fallback if formatting is off
-        idiom = clean_response
-        explanation = ""
-
-    return idiom, explanation
-
+pipe = pipeline(task="text-generation", model="bigscience/bloomz-560m")
+    
+def generate_idiom(situation):
+    response = pipe(f"Give me a Chinese idiom for this: {situation}.")
+    print(response)
+    pinyin = ""
+    meaning = ""
+    return response, f"{pinyin}\n\n{meaning}"
 
 with gr.Blocks(css="""
     .idiom-output {
@@ -100,4 +59,4 @@ with gr.Blocks(css="""
 
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()