Update app.py

app.py

@@ -1,8 +1,10 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load the model
-model = pipeline("text-generation", model="karthikqnq/qnqgpt2")
+# Load the model and tokenizer
+model_name = "karthikqnq/qnqgpt2"
+model = AutoModelForCausalLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 def respond(
     message,
@@ -19,17 +21,23 @@ def respond(
             prompt += f"User: {user_msg}\nAssistant: {assistant_msg}\n"
     prompt += f"User: {message}\nAssistant: "
 
+    # Tokenize the input prompt
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
+
     # Generate response
-    response = model(
-        prompt,
+    outputs = model.generate(
+        **inputs,
         max_length=max_tokens,
         temperature=temperature,
         top_p=top_p,
         do_sample=True,
         num_return_sequences=1
-    )[0]['generated_text']
+    )
+
+    # Decode the output and extract only the assistant's response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Extract only the assistant's response
+    # Extract the assistant's reply after "Assistant:"
     try:
         assistant_response = response.split("Assistant: ")[-1].strip()
     except:
@@ -72,4 +80,4 @@ demo = gr.ChatInterface(
 )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()