Update app.py

app.py

@@ -1,12 +1,8 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
 import spaces
 
-import os
-from threading import Thread
-from typing import Iterator
-
 # Define quantization configuration
 quantization_config = BitsAndBytesConfig(
     load_in_4bit=True,  # Specify 4-bit quantization
@@ -18,64 +14,39 @@ quantization_config = BitsAndBytesConfig(
 # Load the tokenizer and quantized model from Hugging Face
 model_name = "llSourcell/medllama2_7b"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
 # Load model with quantization
 model = AutoModelForCausalLM.from_pretrained(model_name, 
                                              quantization_config=quantization_config, 
                                              device_map="auto")
 model.eval()
-max_token_length = 4096
-
-@spaces.GPU(duration=15)
-def generate(
-    message: str,
-    chat_history: list[tuple[str, str]],
-    max_new_tokens: int = 1024,
-    temperature: float = 0.6,
-    top_p: float = 0.9,
-    top_k: int = 50,
-    repetition_penalty: float = 1.2,
-) -> Iterator[str]:
-    conversation = []
-    for user, assistant in chat_history:
-        conversation.extend(
-            [
-                {"role": "user", "content": user},
-                {"role": "assistant", "content": assistant},
-            ]
-        )
-    conversation.append({"role": "user", "content": message})
 
-    input_ids = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, return_tensors="pt")
-    if input_ids.shape[1] > max_token_length:
-        input_ids = input_ids[:, -max_token_length:]
-        gr.Warning(f"Trimmed input from conversation as it was longer than {max_token_length} tokens.")
-    input_ids = input_ids.to(model.device)
+def format_history(msg: str, history: list[list[str, str]], system_prompt: str):
+    chat_history = system_prompt
+    for query, response in history:
+        chat_history += f"\nUser: {query}\nAssistant: {response}"
+    chat_history += f"\nUser: {msg}\nAssistant:"
+    return chat_history
 
-    streamer = TextIteratorStreamer(tokenizer, timeout=20.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,
-        top_p=top_p,
-        top_k=top_k,
-        temperature=temperature,
-        num_beams=1,
-        repetition_penalty=repetition_penalty,
-    )
-    t = Thread(target=model.generate, kwargs=generate_kwargs)
-    t.start()
+@spaces.GPU(duration=30)
+def generate_response(msg: str, history: list[list[str, str]], system_prompt: str):
+    chat_history = format_history(msg, history, system_prompt)
+    
+    # Tokenize the input prompt
+    inputs = tokenizer(chat_history, return_tensors="pt").to("cuda")
+    
+    # Generate a response using the model
+    outputs = model.generate(inputs["input_ids"], max_length=1024, pad_token_id=tokenizer.eos_token_id)
 
-    outputs = []
-    for text in streamer:
-        outputs.append(text)
-        yield "".join(outputs)
+    # Decode the response back to a string
+    response = tokenizer.decode(outputs[:, inputs["input_ids"].shape[-1]:][0], skip_special_tokens=True)
+    
+    # Yield the generated response
+    yield response
 
 # Define the Gradio ChatInterface
 chatbot = gr.ChatInterface(
-                fn=generate,
+                generate_response,
                 chatbot=gr.Chatbot(
                         height="64vh"
                     ),