Update app.py

app.py

@@ -1,17 +1,24 @@
 import gradio as gr
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
+from transformers import AutoTokenizer
+from auto_gptq import AutoGPTQForCausalLM
 
-tokenizer = AutoTokenizer.from_pretrained("kimnt93/chat-llama2-1b-1.0")
-model = AutoModelForCausalLM.from_pretrained("kimnt93/chat-llama2-1b-1.0")
+model_path = 'TheBloke/SOLAR-10.7B-Instruct-v1.0-uncensored-GPTQ'
+tokenizer_path = 'TheBloke/SOLAR-10.7B-Instruct-v1.0-uncensored-GPTQ'
+
+model = AutoGPTQForCausalLM.from_quantized(
+    model_path,
+    disable_exllama=True,
+    device_map='auto',
+    revision='2bit_128g',
+)
+
+tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, trust_remote_code=True)
 
 def generate_response(prompt):
-    conversation = [{'role': 'user', 'content': prompt}]
-    prompt = tokenizer.apply_chat_template(conversation, tokenizer=False, add_generation_prompt=True)
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, use_cache=True, max_length=4096)
-    outputs_text = tokenizer.decode(outputs[0])
-    return outputs_text
+    input_ids = tokenizer(prompt, return_tensors='pt').input_ids.to('cuda')
+    outputs = model.generate(input_ids=input_ids, max_length=128)
+    decoded_output = tokenizer.decode(outputs[0])
+    return decoded_output
 
 iface = gr.Interface(fn=generate_response, inputs="text", outputs="text")
 iface.launch()