Update app.py

app.py

@@ -55,7 +55,7 @@ h1 {
 model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.2",use_auth_token=access_token)
 model = PeftModel.from_pretrained(model, "physician-ai/mistral-finetuned1",use_auth_token=access_token)
 tokenizer = AutoTokenizer.from_pretrained("physician-ai/mistral-finetuned1",use_auth_token=access_token)
-text_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512, temperature=0.8, top_p=0.95, repetition_penalty=1.15)
+text_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=1024, temperature=0.8, top_p=0.95, repetition_penalty=1.15)
 
 terminators = [
     tokenizer.eos_token_id,
@@ -63,51 +63,41 @@ terminators = [
 ]
 
 
-@spaces.GPU(duration=120)
-def chat_llama3_8b(message: str, 
-              history: list, 
-              temperature: float, 
-              max_new_tokens: int
-             ) -> str:
-    """
-    Generate a streaming response using the llama3-8b model.
-    Args:
-        message (str): The input message.
-        history (list): The conversation history used by ChatInterface.
-        temperature (float): The temperature for generating the response.
-        max_new_tokens (int): The maximum number of new tokens to generate.
-    Returns:
-        str: The generated response.
-    """
-    conversation = []
-    for user, assistant in history:
-        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
-    conversation.append({"role": "user", "content": message})
-
-    input_ids = tokenizer.apply_chat_template(conversation, return_tensors="pt").to(model.device)
-    
-    streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
-
-    generate_kwargs = dict(
-        input_ids= input_ids,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        do_sample=True,
-        temperature=temperature,
-        eos_token_id=terminators,
-    )
-    # This will enforce greedy generation (do_sample=False) when the temperature is passed 0, avoiding the crash.             
-    if temperature == 0:
-        generate_kwargs['do_sample'] = False
+def generate_response(input_ids, generate_kwargs):
+    try:
+        # Generate the output using the model
+        output = model.generate(**generate_kwargs)
+        return output
+    except Exception as e:
+        print(f"Error during generation: {e}")
+
         
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
-
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        #print(outputs)
-        yield "".join(outputs)
+@gr.cache(max_size=100, expire=3600)
+def chat_llama3_8b(message, history, temperature=0.95, max_new_tokens=512):
+    # Prepare conversation context
+    conversation = [{"role": "user", "content": message}] + [{"role": "assistant", "content": reply} for reply in history]
+    input_ids = tokenizer(conversation, return_tensors="pt", padding=True, truncation=True).input_ids.to(model.device)
+    
+    generate_kwargs = {
+        "input_ids": input_ids,
+        "max_length": input_ids.shape[1] + max_new_tokens,
+        "temperature": temperature,
+        "num_return_sequences": 1
+    }
+
+    # Thread for generating model response
+    output_queue = []
+    response_thread = Thread(target=generate_response, args=(input_ids, generate_kwargs, output_queue))
+    response_thread.start()
+    response_thread.join()  # Wait for the thread to complete
+
+    # Retrieve the output from the queue
+    if output_queue:
+        output = output_queue[0]
+        if output is not None:
+            generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
+            return generated_text
+    return "An error occurred during text generation."
         
 
 # Gradio block