Update app.py

app.py

@@ -1,24 +1,18 @@
 import gradio as gr
+import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load model and tokenizer
-model = AutoModelForCausalLM.from_pretrained("HuggingFaceH4/starchat2-15b-v0.1")
-tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/starchat2-15b-v0.1")
+# Load a smaller model or in half-precision
+model = AutoModelForCausalLM.from_pretrained("distilgpt2", torch_dtype=torch.float16)
+tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
 
-# Function for inference
-def generate_text(input_text):
-    inputs = tokenizer(input_text, return_tensors="pt")
-    outputs = model.generate(**inputs)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+def generate_text(inputs):
+    responses = []
+    for input_text in inputs:
+        input_tensor = tokenizer(input_text, return_tensors="pt")
+        output = model.generate(**input_tensor)
+        responses.append(tokenizer.decode(output[0], skip_special_tokens=True))
+    return responses
 
-# Gradio interface
-iface = gr.Interface(
-    fn=generate_text,
-    inputs="text",
-    outputs="text",
-    title="Project Build",
-    description="Generate text using the StarChat model."
-)
-
-# Launch the app
-iface.launch()
+iface = gr.Interface(fn=generate_text, inputs="text", outputs="text", allow_flagging="never")
+iface.launch()