updates

app.py

@@ -10,16 +10,21 @@ from transformers import (
 MODEL_ID = "jerinaj/lfm-tool-merged"
 
 
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID
 )
 
+
+# Load model
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
     device_map="auto"
 )
 
+model.eval()
+
 
 def predict(messages):
 
@@ -32,21 +37,32 @@ def predict(messages):
     inputs = tokenizer(
         prompt,
         return_tensors="pt"
-    ).to(model.device)
+    )
 
+    inputs = {
+        k: v.to(model.device)
+        for k, v in inputs.items()
+    }
 
-    output = model.generate(
-        **inputs,
-        max_new_tokens=256,
-        temperature=0.1
-    )
 
+    with torch.inference_mode():
 
-    return tokenizer.decode(
-        output[0][inputs.input_ids.shape[1]:],
+        output = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            temperature=0.1,
+            do_sample=True
+        )
+
+
+    response = tokenizer.decode(
+        output[0][inputs["input_ids"].shape[1]:],
         skip_special_tokens=True
     )
 
+    return response
+
+
 
 demo = gr.Interface(
     fn=predict,
@@ -55,4 +71,8 @@ demo = gr.Interface(
 )
 
 
-demo.launch(share=True)
+demo.launch(
+    server_name="0.0.0.0",
+    server_port=7860,
+    share=True
+)