Update app.py

app.py

@@ -2,10 +2,10 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 from peft import PeftModel
 
-# Load the fine-tuned model and tokenizer
-base_model = AutoModelForCausalLM.from_pretrained("unsloth/Meta-Llama-3.1-8B-Instruct")  # No bnb-4bit here
+# Use a smaller model to reduce memory usage
+base_model = AutoModelForCausalLM.from_pretrained("meta/llama-2-7b-hf")  # Smaller model
 model = PeftModel.from_pretrained(base_model, "DarkAngel/gitallama")
-tokenizer = AutoTokenizer.from_pretrained("unsloth/Meta-Llama-3.1-8B-Instruct")
+tokenizer = AutoTokenizer.from_pretrained("meta/llama-2-7b-hf")  # Use the tokenizer for the smaller model
 
 def generate_response(shloka, transliteration):
     """
@@ -18,15 +18,13 @@ def generate_response(shloka, transliteration):
         }
     ]
     
-    # Ensure the model uses CPU instead of GPU
     inputs = tokenizer.apply_chat_template(
         input_message,
         tokenize=True,
-        add_generation_prompt=True,  # Enable for generation
+        add_generation_prompt=True,
         return_tensors="pt"
-    ).to("cpu")  # Use CPU
+    ).to("cpu")  # Ensure CPU usage
 
-    # Generate response
     text_streamer = TextStreamer(tokenizer, skip_prompt=True)
     generated_tokens = model.generate(
         input_ids=inputs,
@@ -39,7 +37,6 @@ def generate_response(shloka, transliteration):
 
     raw_response = tokenizer.decode(generated_tokens[0], skip_special_tokens=True)
 
-    # Format the response
     try:
         sections = raw_response.split("Hindi Meaning:")
         english_meaning = sections[0].strip()
@@ -47,14 +44,12 @@ def generate_response(shloka, transliteration):
         hindi_meaning = hindi_and_word[0].strip()
         word_meaning = hindi_and_word[1].strip()
 
-        # Format response for better readability
         formatted_response = (
             f"English Meaning:\n{english_meaning}\n\n"
             f"Hindi Meaning:\n{hindi_meaning}\n\n"
             f"Word Meaning:\n{word_meaning}"
         )
     except IndexError:
-        # In case the response format is not as expected
         formatted_response = raw_response
 
     return formatted_response