updates

app.py

@@ -1,20 +1,58 @@
 import gradio as gr
-from llama_cpp import Llama
-from huggingface_hub import hf_hub_download
+import torch
 
-model_path = hf_hub_download(
-    repo_id="jerinaj/lfm-tool",
-    filename="LFM2-350M.Q4_K_M.gguf"
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM
 )
 
-llm = Llama(
-    model_path=model_path,
-    n_ctx=1024,
-    n_threads=2
+
+MODEL_ID = "jerinaj/lfm-tool-merged"
+
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto"
+)
+
+
+def predict(messages):
+
+    prompt = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True
+    )
+
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt"
+    ).to(model.device)
+
+
+    output = model.generate(
+        **inputs,
+        max_new_tokens=256,
+        temperature=0.1
+    )
+
+
+    return tokenizer.decode(
+        output[0][inputs.input_ids.shape[1]:],
+        skip_special_tokens=True
+    )
+
+
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.JSON(),
+    outputs=gr.Textbox()
 )
 
-def chat(prompt):
-    output = llm(prompt, max_tokens=100)
-    return output["choices"][0]["text"]
 
-gr.Interface(fn=chat, inputs="text", outputs="text").launch()
+demo.launch()

requirements.txt

@@ -1,3 +1,6 @@
-llama-cpp-python==0.2.90
+torch
+transformers
+accelerate
 gradio
-huggingface_hub
+sentencepiece
+protobuf