Update app.py

app.py

@@ -4,7 +4,6 @@ import torch
 import logging
 import sys
 import gc
-from contextlib import contextmanager
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -17,22 +16,25 @@ if torch.cuda.is_available():
 
 try:
     logger.info("Loading tokenizer...")
-    model_id = "htigenai/finetune_test_2"
+    # Use the base model's tokenizer instead
+    base_model_id = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
     tokenizer = AutoTokenizer.from_pretrained(
-        model_id,
-        use_fast=False  # Use slow tokenizer to save memory
+        base_model_id,
+        use_fast=True,
+        trust_remote_code=True
     )
     tokenizer.pad_token = tokenizer.eos_token
     logger.info("Tokenizer loaded successfully")
 
-    logger.info("Loading model in 8-bit...")
+    logger.info("Loading fine-tuned model in 8-bit...")
+    model_id = "htigenai/finetune_test_2"
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         device_map="auto",
-        load_in_8bit=True,  # Load in 8-bit instead of 4-bit
+        load_in_8bit=True,
         torch_dtype=torch.float16,
         low_cpu_mem_usage=True,
-        max_memory={0: "12GB", "cpu": "4GB"}  # Limit memory usage
+        max_memory={0: "12GB", "cpu": "4GB"}
     )
     model.eval()
     logger.info("Model loaded successfully in 8-bit")
@@ -43,16 +45,15 @@ try:
 
     def generate_text(prompt, max_tokens=100, temperature=0.7):
         try:
-            # Format the prompt
+            # Format prompt with chat template
             formatted_prompt = f"### Human: {prompt}\n\n### Assistant:"
             
-            # Generate with memory-efficient settings
             inputs = tokenizer(
                 formatted_prompt,
                 return_tensors="pt",
                 padding=True,
                 truncation=True,
-                max_length=256  # Limit input length
+                max_length=256
             ).to(model.device)
 
             with torch.inference_mode():
@@ -72,11 +73,11 @@ try:
 
             response = tokenizer.decode(outputs[0], skip_special_tokens=True)
             
-            # Extract only the assistant's response
+            # Extract assistant's response
             if "### Assistant:" in response:
                 response = response.split("### Assistant:")[-1].strip()
             
-            # Clean up memory after generation
+            # Clean up
             del outputs, inputs
             gc.collect()
             torch.cuda.empty_cache()
@@ -87,7 +88,7 @@ try:
             logger.error(f"Error during generation: {str(e)}")
             return f"Error generating response: {str(e)}"
 
-    # Create a more memory-efficient Gradio interface
+    # Create Gradio interface
     iface = gr.Interface(
         fn=generate_text,
         inputs=[
@@ -117,7 +118,7 @@ try:
             lines=5
         ),
         title="HTIGENAI Reflection Analyzer (8-bit)",
-        description="8-bit quantized text generation. Please keep prompts concise for best results.",
+        description="Using Llama 3.1 base tokenizer with fine-tuned model. Keep prompts concise for best results.",
         examples=[
             ["What is machine learning?", 50, 0.7],
             ["Explain quantum computing", 50, 0.7],
@@ -125,7 +126,7 @@ try:
         cache_examples=False
     )
 
-    # Launch with minimal memory usage
+    # Launch interface
     iface.launch(
         server_name="0.0.0.0",
         share=False,