Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import torch
 from flask import Flask, jsonify, request
 from flask_cors import CORS 
-from transformers import GPTNeoForCausalLM, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments
+from transformers import GPTNeoForCausalLM, AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, TrainingArguments, GenerationConfig 
 # Set the HF_HOME environment variable to a writable directory
 os.environ["HF_HOME"] = "/workspace/huggingface_cache"  # Change this to a writable path in your space
 
@@ -63,6 +63,8 @@ def generate_response(user_input, model_id):
 
     
     inputs = tokenizer(prompt, return_tensors="pt").to('cuda')
+    
+    outputs = model.generate(**inputs, generation_config=generation_config)
     response = (tokenizer.decode(outputs[0], skip_special_tokens=True)) 
      
     cleaned_response = response.replace("User:", "").replace("Assistant:", "").strip()