modified app.py 3

README.md

@@ -3,27 +3,64 @@ title: LLaMA 7B Server
 emoji: 🤖
 colorFrom: blue
 colorTo: purple
-sdk: gradio
-sdk_version: "4.17.0"
-app_file: app.py
+sdk: fastapi
+sdk_version: "0.95.0"
+app_file: main.py
 pinned: false
 ---
 
 # LLaMA 7B Server
 
-A web interface for interacting with the LLaMA 7B model.
+A FastAPI-based server for interacting with the LLaMA 7B model.
 
 ## Features
 
 - [x] Text generation
-- [x] Chat interface
 - [x] Model parameters configuration
+- [x] REST API interface
 
-## How to Use
+## API Usage
 
-1. Enter your prompt in the text box
-2. Click "Generate" or press Enter
-3. View the model's response below
+### Text Generation
+
+Make a POST request to `/generate` with the following JSON body:
+
+```json
+{
+    "prompt": "your prompt here",
+    "max_length": 2048,
+    "num_beams": 3,
+    "early_stopping": true,
+    "no_repeat_ngram_size": 3
+}
+```
+
+Example using curl:
+
+```bash
+curl -X POST http://localhost:7860/generate \
+  -H "Content-Type: application/json" \
+  -d '{"prompt": "Hello, how are you?"}'
+```
+
+Example using Python:
+
+```python
+import requests
+
+url = "http://localhost:7860/generate"
+data = {
+    "prompt": "Hello, how are you?",
+    "max_length": 2048,
+    "num_beams": 3,
+    "early_stopping": True,
+    "no_repeat_ngram_size": 3
+}
+
+response = requests.post(url, json=data)
+result = response.json()
+print(result["generated_text"])  # This will contain your generated text
+```
 
 ## Model Details
 

app.py

@@ -1,39 +1,38 @@
-import gradio as gr
+from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
+app = FastAPI()
+
 # Load model and tokenizer once on startup
 tokenizer = AutoTokenizer.from_pretrained("Salesforce/codet5p-220m")
 model = AutoModelForSeq2SeqLM.from_pretrained("Salesforce/codet5p-220m")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = model.to(device)
 
-def generate(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+class GenerationRequest(BaseModel):
+    prompt: str
+    max_length: int = 2048
+    num_beams: int = 3
+    early_stopping: bool = True
+    no_repeat_ngram_size: int = 3
+
+@app.post("/generate")
+async def generate_text(request: GenerationRequest):
+    inputs = tokenizer(request.prompt, return_tensors="pt").to(device)
     outputs = model.generate(
         **inputs,
-        max_length=2048,
-        num_beams=3,
-        early_stopping=True,
-        no_repeat_ngram_size=3,
+        max_length=request.max_length,
+        num_beams=request.num_beams,
+        early_stopping=request.early_stopping,
+        no_repeat_ngram_size=request.no_repeat_ngram_size,
         eos_token_id=tokenizer.eos_token_id,
         pad_token_id=tokenizer.pad_token_id,
     )
     output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return output_text
-
-# Create Gradio interface
-iface = gr.Interface(
-    fn=generate,
-    inputs=gr.Textbox(lines=10, label="Input Prompt"),
-    outputs=gr.Textbox(label="Generated Output"),
-    title="LLaMA 7B Server",
-    description="A web interface for interacting with the LLaMA 7B model.",
-    allow_flagging="never",
-    allow_api=True,
-    queue=False
-)
+    return {"generated_text": output_text}
 
-# Launch the interface
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

requirements.txt

@@ -2,4 +2,3 @@ fastapi
 uvicorn[standard]
 transformers
 torch
-gradio>=4.17.0

start.sh

@@ -1,3 +1,2 @@
 #!/bin/bash
-python app.py
 uvicorn app:app --host 0.0.0.0 --port 7860