Update app.py

Optimized inference with FA and KV cache.

app.py

@@ -10,8 +10,17 @@ device = "cuda:0" if torch.cuda.is_available() else "cpu"
 
 # Try loading the model with explicit error handling
 try:
-    veri_model = AutoModelForCausalLM.from_pretrained(veri_model_path, device_map="auto", torch_dtype="auto")
     veri_tokenizer = AutoTokenizer.from_pretrained(veri_model_path)
+
+    veri_model = AutoModelForCausalLM.from_pretrained(
+        veri_model_path, 
+        device_map="auto", 
+        torch_dtype="auto",
+        trust_remote_code=True,
+        use_cache=True,  # Enable KV caching
+        attn_implementation="flash_attention_2" if torch.cuda.is_available() else None
+    )
+
 except Exception as e:
     print(f"Model loading error: {e}")
     veri_model = None
@@ -26,27 +35,46 @@ def generate_response(user_message, history):
         return history
         
     # Simple generation without streaming first
-    system_message = "You are VeriThoughts, a helpful assistant that thinks step by step. You are finetuned from a Qwen model, created by Alibaba Cloud, to answer Verilog coding questions. Make sure your input and output interface has the same names as described in the question. Please start your Verilog code with CODE BEGIN and end with CODE END."
+    system_message = "You are VeriThoughts, a helpful assistant that thinks step by step. You are finetuned from a Qwen model, created by Alibaba Cloud. If you are asked a Verilog question, make sure your input and output interface has the same names as described in the question. If you are asked to generate code, please start your Verilog code with CODE BEGIN and end with CODE END."
     
-    # Create conversation history
     conversation = f"System: {system_message}\n"
-    for h in history:
+    recent_history = history[-3:] if len(history) > 3 else history
+  
+    for h in recent_history:
         conversation += f"User: {h[0]}\nAssistant: {h[1]}\n"
     conversation += f"User: {user_message}\nAssistant:"
     
-    inputs = veri_tokenizer(conversation, return_tensors="pt", truncation=True, max_length=2048).to(device)
+    inputs = veri_tokenizer(
+        conversation,                     
+        return_tensors="pt", 
+        truncation=True, 
+        max_length=2048,
+        padding=True,
+        return_attention_mask=True
+    ).to(device)
     
     with torch.no_grad():
         outputs = veri_model.generate(
             **inputs,
             max_new_tokens=4096,
-            temperature=0.6,
+            temperature=0.7,
             top_p=0.95,
             do_sample=True,
+            top_k=50,            # Top-k sampling for efficiency
+ #           pad_token_id=veri_tokenizer.eos_token_id,
+ #           eos_token_id=veri_tokenizer.eos_token_id,
+            use_cache=True,      # Enable KV caching for faster generation
+            repetition_penalty=1.1,  # Reduce repetition
+            length_penalty=1.0,
+            early_stopping=True,     # Stop early when appropriate
+            num_beams=1,            # Greedy search for speed
             pad_token_id=veri_tokenizer.eos_token_id
         )
     
     response = veri_tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
     
     # Return updated history
     return history + [[user_message, response.strip()]]