Update app.py

app.py

@@ -1,60 +1,70 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import login
-from peft import PeftModel
 import torch
+from peft import PeftModel, PeftConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import login
 import os
+import gradio as gr
 
 # Login to Hugging Face Hub
 access_token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
 login(token=access_token)
 
+# Define model details
+peft_model_id = "kuyesu22/sunbird-ug-lang-v1.0-llama-2-7b-hf-lora"  # Update with the correct ID for your fine-tuned Llama 2 model
+config = PeftConfig.from_pretrained(peft_model_id)
+
+# Load base model and tokenizer
+model = AutoModelForCausalLM.from_pretrained(
+    config.base_model_name_or_path,
+    torch_dtype=torch.float16,    # Mixed precision for faster inference
+    device_map="auto",            # Automatically allocate to available devices
+    offload_folder="./offload"    # Directory for offloading layers if needed
+)
+tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
 
-# Load model and tokenizer from the Hugging Face Hub
-model_name = "kuyesu22/sunbird-ug-lang-v1.0-llama-2-7b-hf-lora"
-tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
-model = PeftModel.from_pretrained(model, model_name)
+# Load the LoRA fine-tuned model
+model = PeftModel.from_pretrained(model, peft_model_id)
 
-# Ensure the model is in evaluation mode
+# Set model to evaluation mode
 model.eval()
 
-# Define the translation function
-def translate(text, source_lang="Runyankole", target_lang="English"):
-    prompt = f"Translate from {source_lang} to {target_lang}: {text}"
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    
+# Define the inference function for translation
+def make_inference(english_text):
+    # Format the prompt based on the language pair
+    prompt = f"### English:\n{english_text}\n\n### Runyankole:"
+    batch = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(model.device)
+
+    # Generate the translation
     with torch.no_grad():
-        outputs = model.generate(
-            inputs["input_ids"],
-            max_length=100,
-            num_beams=5,
-            early_stopping=True
-        )
-    
-    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return translation
-
-# Set up the Gradio interface
-def runyankole_to_english(text):
-    return translate(text, source_lang="Runyankole", target_lang="English")
-
-def english_to_runyankole(text):
-    return translate(text, source_lang="English", target_lang="Runyankole")
-
-# Create Gradio inputs and interface
-with gr.Blocks() as demo:
-    gr.Markdown("# Runyankole-English Translation Model")
-    
-    with gr.Tab("Runyankole to English"):
-        runyankole_input = gr.Textbox(label="Enter Runyankole Text")
-        english_output = gr.Textbox(label="English Translation")
-        gr.Button("Translate").click(runyankole_to_english, inputs=runyankole_input, outputs=english_output)
-    
-    with gr.Tab("English to Runyankole"):
-        english_input = gr.components.Textbox(label="Enter English Text")
-        runyankole_output = gr.components.Textbox(label="Runyankole Translation")
-        gr.Button("Translate").click(english_to_runyankole, inputs=english_input, outputs=runyankole_output)
-
-# Launch the app
-demo.launch()
+        with torch.cuda.amp.autocast():  # Mixed precision inference for speed
+            output_tokens = model.generate(
+                input_ids=batch["input_ids"],
+                attention_mask=batch["attention_mask"],
+                max_new_tokens=100,
+                do_sample=True,
+                temperature=0.7,
+                num_return_sequences=1,
+                pad_token_id=tokenizer.eos_token_id
+            )
+
+    # Decode the generated tokens to obtain the translation
+    translated_text = tokenizer.decode(output_tokens[0], skip_special_tokens=True)
+    return translated_text
+
+# Gradio Interface
+def launch_gradio_interface():
+    inputs = gr.components.Textbox(lines=2, label="English Text")  # Input text in English
+    outputs = gr.components.Textbox(label="Translated Runyankole Text")  # Output in Runyankole
+
+    # Launch Gradio app
+    gr.Interface(
+        fn=make_inference,
+        inputs=inputs,
+        outputs=outputs,
+        title="Sunbird UG Lang Translator",
+        description="Translate English to Runyankole using Llama 2 model fine-tuned with LoRA.",
+    ).launch()
+
+# Entry point to run the Gradio app
+if __name__ == "__main__":
+    launch_gradio_interface()