Update app.py

app.py

@@ -1,18 +1,18 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load a smaller model or in half-precision
-model = AutoModelForCausalLM.from_pretrained("distilgpt2", torch_dtype=torch.float16)
-tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+# Load the model
+model_name = "HuggingFaceH4/starchat2-15b-v0.1"  # Your main model
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16)
 
-def generate_text(inputs):
-    responses = []
-    for input_text in inputs:
-        input_tensor = tokenizer(input_text, return_tensors="pt", clean_up_tokenization_spaces=True)
-        output = model.generate(**input_tensor)
-        responses.append(tokenizer.decode(output[0], skip_special_tokens=True))
-    return responses
+def generate_text(input_text):
+    # Ensure input is in the correct format
+    input_tensor = tokenizer(input_text, return_tensors="pt", clean_up_tokenization_spaces=True)
+    output = model.generate(**input_tensor)
+    response = tokenizer.decode(output[0], skip_special_tokens=True)
+    return response
 
 iface = gr.Interface(fn=generate_text, inputs="text", outputs="text", allow_flagging="never")
 iface.launch(server_name="0.0.0.0", server_port=7860)