Update app.py

app.py

@@ -14,15 +14,22 @@ async def root():
 @app.post("/hello/")
 def say_hello(msg: Annotated[str, Form()]):
     print("model")
-    tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
-    model = AutoModelForCausalLM.from_pretrained(
-        "google/gemma-2b-it",
-        device_map="cpu",
-        torch_dtype=torch.bfloat16
-    )
-    print("token & msg")
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    checkpoint = "HuggingFaceTB/SmolLM-1.7B-Instruct"
+
+    device = "cpu" # for GPU usage or "cpu" for CPU usage
+    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+    # for multiple GPUs install accelerate and do `model = AutoModelForCausalLM.from_pretrained(checkpoint, device_map="auto")`
+    model = AutoModelForCausalLM.from_pretrained(checkpoint).to(device)
+
+    messages = [{"role": "user", "content": "things about elasticsearch"}]
+    input_text=tokenizer.apply_chat_template(messages, tokenize=False)
+    print(input_text)
+    
     input_ids = tokenizer(msg, return_tensors="pt").to("cpu")
     print("output")
-    outputs = model.generate(**input_ids, max_length=500) 
+    inputs = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    outputs = model.generate(inputs, max_new_tokens=256, temperature=0.6, top_p=0.92, do_sample=True)
+
     print("complete")
     return {"message": tokenizer.decode(outputs[0])}