Switched to using llama_cpp

app.py

@@ -2,12 +2,25 @@ import gradio as gr
 from transformers import pipeline
 import logging
 import re
+from llama_cpp import Llama
+from huggingface_hub import hf_hub_download
+import sys
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
-logging.getLogger('transformers').setLevel(logging.INFO)
 
-llama = pipeline("text-generation", model="filipealmeida/open-llama-3b-v2-pii-transform", model_kwargs={"load_in_8bit": True})
+
+model_repo="filipealmeida/open-llama-3b-v2-pii-transform"
+model_filename="ggml-model-f16.gguf"
+
+def download_model():
+    print("Downloading model...")
+    sys.stdout.flush()
+    file = hf_hub_download(
+            repo_id=model_repo, filename=model_filename
+    )
+    print("Downloaded " + file)
+    return file
 
 def generate_text(prompt, example):  
     logging.debug(f"Received prompt: {prompt}")
@@ -19,21 +32,18 @@ def generate_text(prompt, example):
 
     logging.info(f"Input : {input}")
 
-    output = llama(input, max_length=70)
-    generated_text = output[0]["generated_text"]
+    output = llm(input, max_tokens=200, stop=["\n"])
+    print(output)
+    generated_text = output['choices'][0]['text']
     logging.info(f"Generated text: {generated_text}")
 
-    match = re.search("### Response:\n(.*?)\n", generated_text, re.DOTALL)
-
-    parsed_text = "ERROR"
-    if match:
-        parsed_text = match.group(1).strip()
-    else:
-        print("No matching section found.")
+    parsed_text = generated_text.split("\n")[0]
 
     logging.info(f"Parsed text: {parsed_text}")
     return parsed_text
 
+model = download_model()
+llm = Llama(model_path=model)
 
 # Create a Gradio interface
 interface = gr.Interface(

requirements.txt

@@ -1,4 +1 @@
-transformers==4.31.0
-torch==2.0.1
-bitsandbytes
-accelerate
+llama-cpp-python