Update app.py

app.py

@@ -1,15 +1,14 @@
 import gradio as gr
-import requests
 import os
+from llama_cpp import Llama
 
-API_URL = os.environ.get("API_URL")
+# Load the model
 
-# Hugging Face API configuration
-headers = {"Authorization": f"Bearer {os.environ['HF_TOKEN']}"}
-
-def query(payload):
-    response = requests.post(API_URL, headers=headers, json=payload)
-    return response.json()
+llm = Llama.from_pretrained(
+    repo_id="AndreasThinks/mistral-7b-english-welsh-translate-GGUF",
+    filename="*q4_k_m.gguf",
+    verbose=False
+)
 
 def translate(text, source_lang, target_lang):
     if source_lang == target_lang:
@@ -19,50 +18,30 @@ def translate(text, source_lang, target_lang):
     
     input_text = f"""
     ### Instruction: {instruction}
-
     ### Input: {text}
-
     ### Response:
-"""
+    """
     
-    output = query({
-        "inputs": input_text,
-        "parameters": {
-            "max_new_tokens": 6000,
-            "return_text": False,
-            "return_full_text": False,
-            "handle_long_generation": "hole"
-        }
-    })
-
-    print(output)
+    output = llm(input_text, max_tokens=6000, stop=["### Input:", "### Instruction:"], echo=True)
     
-    translated_text = output[0]['generated_text']
+    translated_text = output['choices'][0]['text'].split("### Response:")[-1].strip()
     
     return translated_text, input_text, source_lang, target_lang
 
 def continue_generation(translated_text, input_text, source_lang, target_lang):
     full_text = f"{input_text}{translated_text}"
     
-    output = query({
-        "inputs": full_text,
-        "parameters": {
-            "max_new_tokens": 8000,
-            "return_text": False,
-            "return_full_text": False,
-            "handle_long_generation": "hole"
-        }
-    })
+    output = llm(full_text, max_tokens=8000, stop=["### Input:", "### Instruction:"], echo=True)
     
-    new_translated_text = output[0]['generated_text']
-    updated_translated_text = translated_text + new_translated_text
+    new_translated_text = output['choices'][0]['text'].split("### Response:")[-1].strip()
+    updated_translated_text = translated_text + " " + new_translated_text
     
     return updated_translated_text, input_text, source_lang, target_lang
 
 # Create the Gradio interface
 with gr.Blocks() as iface:
     gr.Markdown("# English-Welsh Translator")
-    gr.Markdown("Translate text between English and Welsh using a Hugging Face Inference Endpoint.")
+    gr.Markdown("Translate text between English and Welsh using a local LLM with llama-cpp.")
     
     with gr.Row():
         input_text = gr.Textbox(label="Enter text to translate")