Update app.py

app.py

@@ -1,16 +1,40 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from transformers import BitsAndBytesConfig
+import gc
+
+# Configure 8-bit quantization
+quantization_config = BitsAndBytesConfig(
+    load_in_8bit=True,
+    llm_int8_threshold=6.0,
+    llm_int8_has_fp16_weight=False,
+)
 
 # Load model and tokenizer
 model_name = "Spestly/AwA-1.5B"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32, low_cpu_mem_usage=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    quantization_config=quantization_config,
+    low_cpu_mem_usage=True,
+    torch_dtype=torch.float32,
+)
 
-# Set to evaluation mode
+# Optimizations
+model.config.use_cache = True
+torch.backends.cudnn.benchmark = False
+torch._C._jit_set_profiling_executor(False)
 model.eval()
 
+# Clear memory
+gc.collect()
+torch.cuda.empty_cache() if torch.cuda.is_available() else None
+
 def generate_response(message, history):
+    gc.collect()
+    
     instruction = (
         "You are an LLM called AwA. Aayan Mishra finetunes you. Anthropic does NOT train you. "
         "You are a Qwen 2.5 fine-tune. Your purpose is the help the user accomplish their request to the best of your abilities. "
@@ -18,38 +42,77 @@ def generate_response(message, history):
         f"### Instruction:\n{message}\n\n### Response:"
     )
     
-    inputs = tokenizer(instruction, return_tensors="pt")
-    
-    streamed_output = ""
+    inputs = tokenizer(
+        instruction,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+        max_length=512
+    )
     
-    # Generate with streaming
-    with torch.no_grad():
-        for output in model.generate(
-            **inputs,
-            max_new_tokens=600,
-            num_return_sequences=1,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id,
-            streaming=True,  # Enable streaming
-            yield_per_token=True  # Yield after each token
-        ):
-            next_token = tokenizer.decode(output[-1], skip_special_tokens=True)
-            streamed_output += next_token
-            yield streamed_output.split("### Response:")[-1].strip()
+    try:
+        # Generate initial sequence
+        generated_ids = []
+        past_key_values = None
+        attention_mask = inputs["attention_mask"]
+        
+        with torch.no_grad():
+            for _ in range(400):  # max_new_tokens
+                outputs = model(
+                    input_ids=inputs["input_ids"] if not generated_ids else torch.tensor([[token] for token in generated_ids[-1:]], device=model.device),
+                    attention_mask=attention_mask,
+                    past_key_values=past_key_values,
+                    use_cache=True,
+                )
+                
+                next_token_logits = outputs.logits[:, -1, :]
+                past_key_values = outputs.past_key_values
+                
+                # Apply temperature and top-p sampling
+                probs = torch.nn.functional.softmax(next_token_logits / 0.7, dim=-1)
+                sorted_probs, sorted_indices = torch.sort(probs, descending=True)
+                cumsum_probs = torch.cumsum(sorted_probs, dim=-1)
+                idx_to_remove = cumsum_probs > 0.9
+                idx_to_remove[:, 1:] = idx_to_remove[:, :-1].clone()
+                idx_to_remove[:, 0] = 0
+                sorted_probs[idx_to_remove] = 0.0
+                sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)
+                next_token = torch.multinomial(sorted_probs, num_samples=1)
+                next_token = sorted_indices.gather(-1, next_token)
+                
+                generated_ids.append(next_token.item())
+                attention_mask = torch.cat([attention_mask, torch.ones((1, 1), device=model.device)], dim=-1)
+                
+                # Decode the current token and yield
+                current_text = tokenizer.decode(generated_ids, skip_special_tokens=True)
+                if "### Response:" in current_text:
+                    response_text = current_text.split("### Response:")[-1].strip()
+                    yield response_text
+                
+                # Check for end of generation
+                if next_token.item() == tokenizer.eos_token_id:
+                    break
+                
+    except Exception as e:
+        yield f"An error occurred: {str(e)}"
+    finally:
+        gc.collect()
+        torch.cuda.empty_cache() if torch.cuda.is_available() else None
 
-# Create the Gradio interface with streaming enabled
+# Create Gradio interface
 iface = gr.ChatInterface(
     generate_response,
-    chatbot=gr.Chatbot(height=600, type="messages"),
+    chatbot=gr.Chatbot(
+        height=400,
+        type="messages",
+    ),
     textbox=gr.Textbox(
         placeholder="Type your message here...",
         container=False,
         scale=7
     ),
-    title="AwA-1.5B 🔎 - Experimental",
-    description="Chat with AwA (Answers with Athena). Please note that since AwA is an experimental model, some outputs may not be accurate/expected!",
+    title="AwA-1.5B 🔎 - CPU Optimized",
+    description="Chat with AwA (Answers with Athena). Optimized for CPU operation.",
     theme="ocean",
     examples=[
         "How can CRISPR help us Humans?",
@@ -59,5 +122,11 @@ iface = gr.ChatInterface(
     type="messages"
 )
 
-iface.queue()  # Enable queuing for streaming
-iface.launch()
+iface.queue(max_size=5)
+iface.launch(
+    share=False,
+    debug=False,
+    show_error=True,
+    server_name="0.0.0.0",
+    server_port=7860,
+)