Update app.py

app.py

@@ -1,27 +1,34 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
-# Load the model
-model_id = "HuggingFaceH4/zephyr-7b-alpha"  # or another instruct-tuned model
+# Load model and tokenizer (use public model)
+model_id = "HuggingFaceH4/zephyr-7b-alpha"  # or openchat/openchat-3.5
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype="auto")
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
 
-# AI Function
-def generate_json(prompt):
-    system_prompt = "You are an AI that generates valid and clean JSON objects based on descriptions."
-    input_prompt = f"{system_prompt}\n\nDescription:\n{prompt}\n\nJSON:\n"
-    output = generator(input_prompt, max_new_tokens=256, do_sample=False)[0]["generated_text"]
+# Chat memory
+chat_history = []
 
-    # Remove original prompt from generated text
-    json_part = output.replace(input_prompt, "").strip()
-    return json_part
+# Main chat function
+def chat(user_input, history):
+    prompt = ""
+    for i, (user, bot) in enumerate(history):
+        prompt += f"User: {user}\nAssistant: {bot}\n"
+    prompt += f"User: {user_input}\nAssistant:"
 
-# Gradio UI
-gr.Interface(
-    fn=generate_json,
-    inputs=gr.Textbox(label="Describe the JSON you want", lines=4),
-    outputs=gr.Textbox(label="Generated JSON", lines=20),
-    title="🧠 JSON Generator with LLM",
-    description="Enter a plain language description and get a structured JSON output."
-).launch()
+    response = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7)[0]["generated_text"]
+    answer = response.split("Assistant:")[-1].strip()
+
+    history.append((user_input, answer))
+    return history, history
+
+# Interface
+chatbot = gr.Chatbot()
+demo = gr.Interface(fn=chat, inputs=["text", "state"], outputs=[chatbot, "state"], 
+                    title="ChatBot", 
+                    description="Normal chatbot like ChatGPT")
+
+if __name__ == "__main__":
+    demo.launch()