Update app.py

app.py

@@ -1,41 +1,14 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
 
-DEFAULT_SYSTEM_PROMPT = "あなたは誠実で優秀な日本人のアシスタントです。特に指示が無い場合は、常に日本語で回答してください。"
-text = "仕事の熱意を取り戻すためのアイデアを5つ挙げてください。"
+tokenizer = AutoTokenizer.from_pretrained("lightblue/ao-karasu-72B")
+model = AutoModelForCausalLM.from_pretrained("lightblue/ao-karasu-72B", device_map="auto")
 
-model_name = "elyza/Llama-3-ELYZA-JP-8B"
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype="auto",
-    device_map="auto",
-)
-model.eval()
+messages = [{"role": "system", "content": "あなたはAIアシスタントです。"}]
+messages.append({"role": "user", "content": "イギリスの首相は誰ですか？"})
 
-messages = [
-    {"role": "system", "content": DEFAULT_SYSTEM_PROMPT},
-    {"role": "user", "content": text},
-]
-prompt = tokenizer.apply_chat_template(
-    messages,
-    tokenize=False,
-    add_generation_prompt=True
-)
-token_ids = tokenizer.encode(
-    prompt, add_special_tokens=False, return_tensors="pt"
-)
+prompt = tokenizer.apply_chat_template(conversation=messages, add_generation_prompt=True, tokenize=False)
 
-with torch.no_grad():
-    output_ids = model.generate(
-        token_ids.to(model.device),
-        max_new_tokens=1200,
-        do_sample=True,
-        temperature=0.6,
-        top_p=0.9,
-    )
-output = tokenizer.decode(
-    output_ids.tolist()[0][token_ids.size(1):], skip_special_tokens=True
-)
-print(output)
+pipe(prompt, max_new_tokens=100, do_sample=False, temperature=0.0, return_full_text=False)