refactor

app.py

@@ -34,7 +34,7 @@ interface = gr.Interface(
     fn=metadata_transformer.translate,  # This function will handle the logic and transformations.
     inputs=inputs,
     outputs="text",
-    # live=True,
+    live=True,
     title="Schema Transformer with Llama2",
     description="Paste your schema, specify the target, and get the transformed schema."
 )

metadata_transformer.py

@@ -1,34 +1,29 @@
-from transformers import pipeline
-from transformers import AutoTokenizer
-import torch
 import os
 from huggingface_hub import login
+from easyllm.clients import huggingface
+from easyllm.prompt_utils import build_llama2_prompt
 
 TOKEN = os.environ.get("TOKEN")
 
 login(token=TOKEN)
 
-def translate(schema_input, schema_target):
-
-
-    model = "meta-llama/Llama-2-7b-chat-hf"
-
-    tokenizer = AutoTokenizer.from_pretrained(model)
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        torch_dtype=torch.float16,
-        device_map="auto",
+huggingface.prompt_builder = build_llama2_prompt
+
+system_message = """
+You are a metadata schema translator. You translate metadata from one schema to another.
+"""
+
+def translate(inputs):
+    schema_input, schema_target = inputs
+    propmpt = '{} \n Translate the schema metadata file above to the schema: {}'.format(schema_input, schema_target),
+    response = huggingface.ChatCompletion.create(
+        model="meta-llama/Llama-2-70b-chat-hf",
+        messages=[
+            {"role": "system", "content": system_message},
+            {"role": "user", "content": propmpt},
+        ],
+          temperature=0.9,
+          top_p=0.6,
+          max_tokens=256,
     )
-
-    sequences = pipe(
-        '{} \n Translate the schema metadata file above to the schema: {}'.format(schema_input, schema_target),
-        do_sample=True,
-        top_k=10,
-        num_return_sequences=1,
-        eos_token_id=tokenizer.eos_token_id,
-        max_length=200,
-    )
-    return sequences[0]['generated_text']
-
-
+    return response['choices'][0]['message']['content']

requirements.txt

@@ -1,3 +1 @@
-torch==2.0.0
-transformers==4.34.0
-accelerate
+easyllm