Upload app.py with huggingface_hub

app.py

@@ -1,67 +1,50 @@
 import gradio as gr
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
-import torch
+from huggingface_hub import InferenceClient
+import os
 
-# Model: Mistral 7B Instruct - powerful open model
-MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.2"
+# Use Inference API - no need to load model locally
+MODEL_NAME = "mistralai/Mistral-7B-Instruct-v0.3"
+client = InferenceClient(token=os.environ.get("HF_TOKEN"))
 
-print(f"Loading model: {MODEL_NAME}...")
-
-# Initialize text generation pipeline
-try:
-    generator = pipeline(
-        "text-generation",
-        model=MODEL_NAME,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto" if torch.cuda.is_available() else None,
-        trust_remote_code=True
-    )
-    print("Model loaded successfully!")
-except Exception as e:
-    print(f"Error loading model: {e}")
-    # Fallback to a smaller model
-    MODEL_NAME = "microsoft/DialoGPT-medium"
-    generator = pipeline("text-generation", model=MODEL_NAME)
-    print(f"Loaded fallback model: {MODEL_NAME}")
+print(f"Using model: {MODEL_NAME} via Inference API")
 
 def chat_with_ai(message, history):
-    """Chat with the AI model."""
+    """Chat with the AI model via Inference API."""
     if not message.strip():
         return history
     
-    # Build conversation prompt
-    conversation = ""
+    # Convert history to messages format
+    messages = []
     for user_msg, assistant_msg in history:
-        conversation += f"<|user|>\n{user_msg}</s>\n<|assistant|\n{assistant_msg}</s>\n"
-    conversation += f"<|user|>\n{message}</s>\n<|assistant|"
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": assistant_msg})
+    messages.append({"role": "user", "content": message})
     
     try:
-        # Generate response
-        response = generator(
-            conversation,
-            max_new_tokens=512,
+        # Call the Inference API
+        response = client.chat.completions.create(
+            model=MODEL_NAME,
+            messages=messages,
+            max_tokens=512,
             temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=generator.tokenizer.eos_token_id
+            top_p=0.9
         )
         
-        # Extract only the new response
-        full_text = response[0]['generated_text']
-        # Get the part after the last assistant tag
-        if "<|assistant| " in full_text:
-            assistant_response = full_text.split("<|assistant|")[-1].strip()
-        else:
-            assistant_response = full_text[len(conversation):].strip()
-        
-        # Clean up any remaining tags
-        assistant_response = assistant_response.replace("</s>", "").strip()
+        assistant_response = response.choices[0].message.content
         
-        if not assistant_response:
-            assistant_response = "I'm thinking... could you ask that again?"
-            
     except Exception as e:
-        assistant_response = f"Sorry, I encountered an error: {str(e)}"
+        # Fallback to text generation if chat fails
+        try:
+            prompt = f"User: {message}\nAssistant:"
+            response = client.text_generation(
+                model=MODEL_NAME,
+                prompt=prompt,
+                max_new_tokens=256,
+                temperature=0.7
+            )
+            assistant_response = response
+        except Exception as e2:
+            assistant_response = f"Sorry, couldn't connect to the model. Error: {str(e)}"
     
     history.append((message, assistant_response))
     return history
@@ -78,10 +61,13 @@ with gr.Blocks(
         max-width: 800px !important;
         margin: auto !important;
     }
+    footer {
+        display: none !important;
+    }
     """
 ) as demo:
     gr.Markdown("# 🤖 AI Chat Assistant")
-    gr.Markdown(f"Powered by **{MODEL_NAME}**")
+    gr.Markdown(f"Powered by **Mistral-7B-Instruct**")
     
     chatbot = gr.Chatbot(
         label="Chat",
@@ -103,6 +89,9 @@ with gr.Blocks(
     with gr.Row():
         clear_btn = gr.Button("Clear Chat", variant="secondary")
     
+    gr.Markdown("---")
+    gr.Markdown("*Space made by: you can already see it*")
+    
     # Event handlers
     msg.submit(chat_with_ai, [msg, chatbot], [chatbot]).then(
         lambda: "", None, [msg]